diff --git a/src/main/java/org/apache/sysds/parser/DMLTranslator.java b/src/main/java/org/apache/sysds/parser/DMLTranslator.java index a8e1667d049..0db739cc901 100644 --- a/src/main/java/org/apache/sysds/parser/DMLTranslator.java +++ b/src/main/java/org/apache/sysds/parser/DMLTranslator.java @@ -1058,6 +1058,7 @@ public void constructHops(StatementBlock sb) { case LIBSVM: case HDF5: case DELTA: + case PARQUET: // columnar/text formats: no block layout (blocksize -1) ae.setOutputParams(ae.getDim1(), ae.getDim2(), ae.getNnz(), ae.getUpdateType(), -1); break; diff --git a/src/main/java/org/apache/sysds/runtime/io/FrameReaderFactory.java b/src/main/java/org/apache/sysds/runtime/io/FrameReaderFactory.java index 5efbf80b83e..da903543f52 100644 --- a/src/main/java/org/apache/sysds/runtime/io/FrameReaderFactory.java +++ b/src/main/java/org/apache/sysds/runtime/io/FrameReaderFactory.java @@ -51,6 +51,8 @@ public static FrameReader createFrameReader(FileFormat fmt, FileFormatProperties case PROTO: // TODO performance improvement: add parallel reader return new FrameReaderProto(); + case PARQUET: + return binaryParallel ? new FrameReaderParquetParallel() : new FrameReaderParquet(); case DELTA: return textParallel ? new FrameReaderDeltaParallel() : new FrameReaderDelta(); default: diff --git a/src/main/java/org/apache/sysds/runtime/io/FrameReaderParquet.java b/src/main/java/org/apache/sysds/runtime/io/FrameReaderParquet.java index ff23e9ea316..79ca2cff38e 100644 --- a/src/main/java/org/apache/sysds/runtime/io/FrameReaderParquet.java +++ b/src/main/java/org/apache/sysds/runtime/io/FrameReaderParquet.java @@ -20,26 +20,41 @@ import java.io.IOException; import java.io.InputStream; +import java.util.Arrays; +import java.util.Comparator; import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.FileSystem; import org.apache.hadoop.fs.Path; -import org.apache.parquet.example.data.Group; +import org.apache.parquet.column.ColumnDescriptor; +import org.apache.parquet.column.ColumnReader; +import org.apache.parquet.column.impl.ColumnReadStoreImpl; +import org.apache.parquet.column.page.PageReadStore; import org.apache.parquet.hadoop.ParquetFileReader; -import org.apache.parquet.hadoop.ParquetReader; -import org.apache.parquet.hadoop.example.GroupReadSupport; -import org.apache.parquet.hadoop.metadata.ParquetMetadata; +import org.apache.parquet.hadoop.metadata.FileMetaData; import org.apache.parquet.hadoop.util.HadoopInputFile; +import org.apache.parquet.io.api.Converter; +import org.apache.parquet.io.api.GroupConverter; +import org.apache.parquet.io.api.PrimitiveConverter; import org.apache.parquet.schema.MessageType; -import org.apache.parquet.schema.PrimitiveType; +import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName; +import org.apache.parquet.schema.Type; +import org.apache.parquet.schema.Type.Repetition; import org.apache.sysds.common.Types.ValueType; import org.apache.sysds.conf.ConfigurationManager; import org.apache.sysds.runtime.DMLRuntimeException; import org.apache.sysds.runtime.frame.data.FrameBlock; +import org.apache.sysds.runtime.frame.data.columns.Array; +import org.apache.sysds.runtime.frame.data.columns.ArrayFactory; import org.apache.sysds.runtime.util.HDFSTool; +import org.apache.sysds.runtime.util.UtilFunctions; /** * Single-threaded frame parquet reader. - * + * + * Decodes through parquet-mr's column API ({@link ColumnReadStoreImpl}/{@link ColumnReader}) directly into + * pre-allocated typed column arrays. The output frame is constructed from the filled arrays without copying. Columns + * whose parquet physical type does not match the requested frame value type are converted per cell instead. */ public class FrameReaderParquet extends FrameReader { @@ -54,104 +69,312 @@ public class FrameReaderParquet extends FrameReader { * @return A FrameBlock containing the data read from the Parquet file. */ @Override - public FrameBlock readFrameFromHDFS(String fname, ValueType[] schema, String[] names, long rlen, long clen) throws IOException, DMLRuntimeException { - // Prepare file access + public FrameBlock readFrameFromHDFS(String fname, ValueType[] schema, String[] names, long rlen, long clen) + throws IOException, DMLRuntimeException { Configuration conf = ConfigurationManager.getCachedJobConf(); Path path = new Path(fname); - - // Check existence and non-empty file - if (!HDFSTool.existsFileOnHDFS(path.toString())) { + if(!HDFSTool.existsFileOnHDFS(path.toString())) throw new IOException("File does not exist on HDFS: " + fname); - } - // Allocate output frame block ValueType[] lschema = createOutputSchema(schema, clen); String[] lnames = createOutputNames(names, clen); - FrameBlock ret = createOutputFrameBlock(lschema, lnames, rlen); - // Read Parquet file - readParquetFrameFromHDFS(path, conf, ret, lschema, rlen, clen); + Object[] dest = new Object[(int) clen]; + for(int c = 0; c < clen; c++) + dest[c] = ArrayFactory.allocateBacking(lschema[c], (int) rlen); + + readParquetFrameFromHDFS(path, conf, dest, lschema, lnames, rlen); - return ret; + // zero-row output stays a schema-only frame (no zero-length column arrays) + if(rlen == 0) + return new FrameBlock(lschema, lnames, 0); + + Array[] columns = new Array[(int) clen]; + for(int c = 0; c < clen; c++) + columns[c] = ArrayFactory.create(lschema[c], dest[c]); + return new FrameBlock(columns, lnames); } /** - * Reads data from a Parquet file on HDFS and fills the provided FrameBlock. - * The method retrieves the Parquet schema from the file footer, maps the required column names - * to their corresponding indices, and then uses a ParquetReader to iterate over each row. - * Data is extracted based on the column type and set into the output FrameBlock. + * Reads an entire Parquet file (or directory of part files) into the pre-allocated column backing arrays. Part + * files, if any, are read sequentially in name order. * - * @param path The HDFS path to the Parquet file. + * @param path The HDFS path to the Parquet file or directory. * @param conf The Hadoop configuration. - * @param dest The FrameBlock to populate with data. - * @param schema The expected value types for the output columns. + * @param dest The per-column backing arrays to populate. + * @param schema The value types of the output columns. + * @param names The names of the output columns. * @param rlen The expected number of rows. - * @param clen The expected number of columns. */ - protected void readParquetFrameFromHDFS(Path path, Configuration conf, FrameBlock dest, ValueType[] schema, long rlen, long clen) throws IOException { - // Retrieve schema from Parquet footer - ParquetMetadata metadata = ParquetFileReader.open(HadoopInputFile.fromPath(path, conf)).getFooter(); - MessageType parquetSchema = metadata.getFileMetaData().getSchema(); - - // Map column names to Parquet schema indices - String[] columnNames = dest.getColumnNames(); - int[] columnIndices = new int[columnNames.length]; - for (int i = 0; i < columnNames.length; i++) { - columnIndices[i] = parquetSchema.getFieldIndex(columnNames[i]); - } + protected void readParquetFrameFromHDFS(Path path, Configuration conf, Object[] dest, ValueType[] schema, + String[] names, long rlen) throws IOException { + FileSystem fs = IOUtilFunctions.getFileSystem(path); + Path[] files = IOUtilFunctions.getSequenceFilePaths(fs, path); + Arrays.sort(files, Comparator.comparing(Path::getName)); - // Read data usind ParquetReader - try (ParquetReader rowReader = ParquetReader.builder(new GroupReadSupport(), path) - .withConf(conf) - .build()) { - - Group group; - int row = 0; - while ((group = rowReader.read()) != null) { - for (int col = 0; col < clen; col++) { - int colIndex = columnIndices[col]; - if (group.getFieldRepetitionCount(colIndex) > 0) { - PrimitiveType.PrimitiveTypeName type = parquetSchema.getType(columnNames[col]).asPrimitiveType().getPrimitiveTypeName(); - switch (type) { - case INT32: - dest.set(row, col, group.getInteger(colIndex, 0)); - break; - case INT64: - dest.set(row, col, group.getLong(colIndex, 0)); - break; - case FLOAT: - dest.set(row, col, group.getFloat(colIndex, 0)); - break; - case DOUBLE: - dest.set(row, col, group.getDouble(colIndex, 0)); - break; - case BOOLEAN: - dest.set(row, col, group.getBoolean(colIndex, 0)); - break; - case BINARY: - dest.set(row, col, group.getBinary(colIndex, 0).toStringUsingUTF8()); - break; - default: - throw new IOException("Unsupported data type: " + type); - } - } else { - dest.set(row, col, null); - } + long off = 0; + for(Path file : files) + off += readSingleParquetFile(file, conf, dest, schema, names, rlen, (int) off); + if(off != rlen) + throw new IOException("Mismatch in row count: expected " + rlen + ", but got " + off); + } + + /** + * Decodes a single Parquet file into the column backing arrays at the given row offset, one row group at a time and + * column-at-a-time within each row group. Thread-safe for distinct row ranges, so the parallel reader assigns each + * file its own offset and shares the output arrays. + * + * @param path The HDFS path to the Parquet file. + * @param conf The Hadoop configuration. + * @param dest The per-column backing arrays to populate. + * @param schema The value types of the output columns. + * @param names The names of the output columns. + * @param rlen The total number of output rows (exclusive upper bound of this file's rows). + * @param rowOffset The row offset of this file's first row. + * @return The number of rows read. + */ + protected int readSingleParquetFile(Path path, Configuration conf, Object[] dest, ValueType[] schema, + String[] names, long rlen, int rowOffset) throws IOException { + final int ncol = schema.length; + try(ParquetFileReader reader = ParquetFileReader.open(HadoopInputFile.fromPath(path, conf))) { + FileMetaData meta = reader.getFooter().getFileMetaData(); + MessageType parquetSchema = meta.getSchema(); + String createdBy = meta.getCreatedBy(); + + // map each requested frame column (by name) to its parquet column descriptor + final ColumnDescriptor[] descs = new ColumnDescriptor[ncol]; + for(int c = 0; c < ncol; c++) + descs[c] = validateDecodable(parquetSchema, names[c]); + + GroupConverter root = dummyConverter(parquetSchema.getFieldCount()); + int off = rowOffset; + PageReadStore pages; + while((pages = reader.readNextRowGroup()) != null) { + int nrow = (int) pages.getRowCount(); + if(off + nrow > rlen) + throw new IOException( + "Mismatch in row count: expected " + rlen + ", but got at least " + (off + nrow)); + ColumnReadStoreImpl store = new ColumnReadStoreImpl(pages, root, parquetSchema, createdBy); + for(int c = 0; c < ncol; c++) { + ColumnReader creader = store.getColumnReader(descs[c]); + int maxDef = descs[c].getMaxDefinitionLevel(); + PrimitiveTypeName ptype = descs[c].getPrimitiveType().getPrimitiveTypeName(); + if(directDecodable(ptype, schema[c])) + decodeColumnInto(creader, maxDef, nrow, ptype, dest[c], off); + else + decodeColumnConvert(creader, maxDef, nrow, ptype, schema[c], dest[c], off); } - row++; + off += nrow; } + return off - rowOffset; + } + } + + /** + * Resolve the descriptor of one parquet column and verify it is decodable: a non-nested primitive of a physical + * type the reader supports (INT96 timestamps and nested/repeated groups are not). + */ + private static ColumnDescriptor validateDecodable(MessageType parquetSchema, String name) throws IOException { + if(!parquetSchema.containsField(name)) + throw new IOException("Column not found in Parquet schema: " + name); + Type t = parquetSchema.getType(name); + if(!t.isPrimitive() || t.isRepetition(Repetition.REPEATED)) + throw new IOException("Nested Parquet columns are not supported: " + name); + PrimitiveTypeName ptype = t.asPrimitiveType().getPrimitiveTypeName(); + switch(ptype) { + case INT32: + case INT64: + case FLOAT: + case DOUBLE: + case BOOLEAN: + case BINARY: + return parquetSchema.getColumnDescription(new String[] {name}); + default: + throw new IOException("Unsupported Parquet type " + ptype + " for column: " + name + + (ptype == PrimitiveTypeName.INT96 ? " (deprecated INT96 timestamps; re-encode as INT64)" : "")); + } + } - // Check frame dimensions - if (row != rlen) { - throw new IOException("Mismatch in row count: expected " + rlen + ", but got " + row); + /** + * Whether the parquet physical type is the one the frame value type's backing array stores, i.e. the column can be + * decoded through the typed direct path without per-value conversion. + */ + private static boolean directDecodable(PrimitiveTypeName ptype, ValueType vt) { + switch(ptype) { + case INT32: + return vt == ValueType.INT32; + case INT64: + return vt == ValueType.INT64; + case FLOAT: + return vt == ValueType.FP32; + case DOUBLE: + return vt == ValueType.FP64; + case BOOLEAN: + return vt == ValueType.BOOLEAN; + default: + return vt == ValueType.STRING; // BINARY + } + } + + /** + * Decode one parquet column of the current row group into a pre-allocated typed array at the given offset. Null + * cells (definition level below max) keep the array default (0 for numerics, null for strings). + */ + private static void decodeColumnInto(ColumnReader creader, int maxDef, int nrow, PrimitiveTypeName ptype, + Object dest, int off) { + final int end = off + nrow; + switch(ptype) { + case INT32: { + int[] a = (int[]) dest; + for(int r = off; r < end; r++) { + if(creader.getCurrentDefinitionLevel() == maxDef) + a[r] = creader.getInteger(); + creader.consume(); + } + break; + } + case INT64: { + long[] a = (long[]) dest; + for(int r = off; r < end; r++) { + if(creader.getCurrentDefinitionLevel() == maxDef) + a[r] = creader.getLong(); + creader.consume(); + } + break; + } + case FLOAT: { + float[] a = (float[]) dest; + for(int r = off; r < end; r++) { + if(creader.getCurrentDefinitionLevel() == maxDef) + a[r] = creader.getFloat(); + creader.consume(); + } + break; + } + case DOUBLE: { + double[] a = (double[]) dest; + for(int r = off; r < end; r++) { + if(creader.getCurrentDefinitionLevel() == maxDef) + a[r] = creader.getDouble(); + creader.consume(); + } + break; + } + case BOOLEAN: { + boolean[] a = (boolean[]) dest; + for(int r = off; r < end; r++) { + if(creader.getCurrentDefinitionLevel() == maxDef) + a[r] = creader.getBoolean(); + creader.consume(); + } + break; + } + default: { // BINARY + String[] a = (String[]) dest; + for(int r = off; r < end; r++) { + if(creader.getCurrentDefinitionLevel() == maxDef) + a[r] = creader.getBinary().toStringUsingUTF8(); + creader.consume(); + } + break; } } } + /** + * Decode one parquet column whose physical type does not match the frame value type, converting each value per cell + * (same semantics as {@link FrameBlock#set(int, int, Object)}, e.g. a DOUBLE file column read into a STRING frame + * column). + */ + private static void decodeColumnConvert(ColumnReader creader, int maxDef, int nrow, PrimitiveTypeName ptype, + ValueType vt, Object dest, int off) { + final int end = off + nrow; + for(int r = off; r < end; r++) { + if(creader.getCurrentDefinitionLevel() == maxDef) + setConverted(dest, vt, r, readValue(creader, ptype)); + creader.consume(); + } + } + + private static Object readValue(ColumnReader creader, PrimitiveTypeName ptype) { + switch(ptype) { + case INT32: + return creader.getInteger(); + case INT64: + return creader.getLong(); + case FLOAT: + return creader.getFloat(); + case DOUBLE: + return creader.getDouble(); + case BOOLEAN: + return creader.getBoolean(); + default: + return creader.getBinary().toStringUsingUTF8(); // BINARY + } + } + + /** + * Store one converted value into the backing array of the given value type. + */ + private static void setConverted(Object dest, ValueType vt, int r, Object val) { + Object converted = UtilFunctions.objectToObject(vt, val); + if(converted == null) + return; + switch(vt) { + case FP64: + ((double[]) dest)[r] = (Double) converted; + break; + case FP32: + ((float[]) dest)[r] = (Float) converted; + break; + case INT64: + case HASH64: + ((long[]) dest)[r] = (Long) converted; + break; + case UINT4: + case UINT8: + case INT32: + case HASH32: + ((int[]) dest)[r] = (Integer) converted; + break; + case BOOLEAN: + ((boolean[]) dest)[r] = (Boolean) converted; + break; + case CHARACTER: + ((char[]) dest)[r] = (Character) converted; + break; + default: + ((String[]) dest)[r] = (String) converted; + break; + } + } + + /** No-op converter tree; the column API requires one, but values are pulled via the typed getters. */ + private static GroupConverter dummyConverter(int nFields) { + final PrimitiveConverter[] leaves = new PrimitiveConverter[nFields]; + for(int i = 0; i < nFields; i++) + leaves[i] = new PrimitiveConverter() { + }; + return new GroupConverter() { + @Override + public Converter getConverter(int fieldIndex) { + return leaves[fieldIndex]; + } + + @Override + public void start() { + } + + @Override + public void end() { + } + }; + } + //not implemented @Override public FrameBlock readFrameFromInputStream(InputStream is, ValueType[] schema, String[] names, long rlen, long clen) throws IOException, DMLRuntimeException { throw new UnsupportedOperationException("Unimplemented method 'readFrameFromInputStream'"); } -} \ No newline at end of file +} diff --git a/src/main/java/org/apache/sysds/runtime/io/FrameReaderParquetParallel.java b/src/main/java/org/apache/sysds/runtime/io/FrameReaderParquetParallel.java index 3d40f53c626..a28a366a317 100644 --- a/src/main/java/org/apache/sysds/runtime/io/FrameReaderParquetParallel.java +++ b/src/main/java/org/apache/sysds/runtime/io/FrameReaderParquetParallel.java @@ -20,6 +20,8 @@ import java.io.IOException; import java.util.ArrayList; +import java.util.Arrays; +import java.util.Comparator; import java.util.List; import java.util.concurrent.Callable; import java.util.concurrent.ExecutorService; @@ -27,91 +29,79 @@ import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.FileSystem; import org.apache.hadoop.fs.Path; -import org.apache.parquet.example.data.Group; -import org.apache.parquet.hadoop.ParquetReader; -import org.apache.parquet.hadoop.example.GroupReadSupport; +import org.apache.parquet.hadoop.ParquetFileReader; +import org.apache.parquet.hadoop.metadata.BlockMetaData; +import org.apache.parquet.hadoop.util.HadoopInputFile; import org.apache.sysds.common.Types.ValueType; import org.apache.sysds.hops.OptimizerUtils; -import org.apache.sysds.runtime.DMLRuntimeException; -import org.apache.sysds.runtime.frame.data.FrameBlock; import org.apache.sysds.runtime.util.CommonThreadPool; /** - * Multi-threaded frame parquet reader. - * + * Multi-threaded frame parquet reader: reads one file per task, each decoding into its own row range of the shared + * column backing arrays. Per-file row offsets are derived from the row counts in the file footers. */ public class FrameReaderParquetParallel extends FrameReaderParquet { - - /** - * Reads a Parquet frame in parallel and populates the provided FrameBlock with the data. - * The method retrieves all file paths from the sequence files at that location, it then determines - * the number of threads to use based on the available files and a configured parallelism setting. - * A thread pool is created to run a reading task for each file concurrently. - * - * @param path The HDFS path to the Parquet file or the directory containing sequence files. - * @param conf The Hadoop configuration. - * @param dest The FrameBlock to be updated with the data read from the files. - * @param schema The expected value types for the frame columns. - * @param rlen The expected number of rows. - * @param clen The expected number of columns. - */ + @Override - protected void readParquetFrameFromHDFS(Path path, Configuration conf, FrameBlock dest, ValueType[] schema, long rlen, long clen) throws IOException, DMLRuntimeException { + protected void readParquetFrameFromHDFS(Path path, Configuration conf, Object[] dest, ValueType[] schema, + String[] names, long rlen) throws IOException { FileSystem fs = IOUtilFunctions.getFileSystem(path); Path[] files = IOUtilFunctions.getSequenceFilePaths(fs, path); + Arrays.sort(files, Comparator.comparing(Path::getName)); int numThreads = Math.min(OptimizerUtils.getParallelBinaryReadParallelism(), files.length); - - // Create and execute read tasks + + long[] offsets = new long[files.length]; + long cumulative = 0; + for(int i = 0; i < files.length; i++) { + offsets[i] = cumulative; + try(ParquetFileReader reader = ParquetFileReader.open(HadoopInputFile.fromPath(files[i], conf))) { + for(BlockMetaData block : reader.getFooter().getBlocks()) + cumulative += block.getRowCount(); + } + } + if(cumulative != rlen) + throw new IOException("Mismatch in row count: expected " + rlen + ", but got " + cumulative); + ExecutorService pool = CommonThreadPool.get(numThreads); try { List tasks = new ArrayList<>(); - for (Path file : files) { - tasks.add(new ReadFileTask(file, conf, dest, schema, clen)); - } + for(int i = 0; i < files.length; i++) + tasks.add(new ReadFileTask(files[i], conf, dest, schema, names, rlen, (int) offsets[i])); - for (Future task : pool.invokeAll(tasks)) { + for(Future task : pool.invokeAll(tasks)) task.get(); - } - } catch (Exception e) { + } + catch(Exception e) { throw new IOException("Failed parallel read of Parquet frame.", e); - } finally { + } + finally { pool.shutdown(); } } private class ReadFileTask implements Callable { - private Path path; - private Configuration conf; - private FrameBlock dest; - @SuppressWarnings("unused") - private ValueType[] schema; - private long clen; + private final Path path; + private final Configuration conf; + private final Object[] dest; + private final ValueType[] schema; + private final String[] names; + private final long rlen; + private final int rowOffset; - public ReadFileTask(Path path, Configuration conf, FrameBlock dest, ValueType[] schema, long clen) { + public ReadFileTask(Path path, Configuration conf, Object[] dest, ValueType[] schema, String[] names, long rlen, + int rowOffset) { this.path = path; this.conf = conf; this.dest = dest; this.schema = schema; - this.clen = clen; + this.names = names; + this.rlen = rlen; + this.rowOffset = rowOffset; } - // When executed, a ParquetReader for the assigned file opens and iterates over each row processing every column. @Override public Object call() throws Exception { - try (ParquetReader reader = ParquetReader.builder(new GroupReadSupport(), path).withConf(conf).build()) { - Group group; - int row = 0; - while ((group = reader.read()) != null) { - for (int col = 0; col < clen; col++) { - if (group.getFieldRepetitionCount(col) > 0) { - dest.set(row, col, group.getValueToString(col, 0)); - } else { - dest.set(row, col, null); - } - } - row++; - } - } + readSingleParquetFile(path, conf, dest, schema, names, rlen, rowOffset); return null; } } diff --git a/src/main/java/org/apache/sysds/runtime/io/FrameWriterFactory.java b/src/main/java/org/apache/sysds/runtime/io/FrameWriterFactory.java index ff38eb395dd..0891ae1397a 100644 --- a/src/main/java/org/apache/sysds/runtime/io/FrameWriterFactory.java +++ b/src/main/java/org/apache/sysds/runtime/io/FrameWriterFactory.java @@ -50,6 +50,8 @@ public static FrameWriter createFrameWriter(FileFormat fmt, FileFormatProperties return binaryParallel ? new FrameWriterBinaryBlockParallel() : new FrameWriterBinaryBlock(); case PROTO: return new FrameWriterProto(); + case PARQUET: + return binaryParallel ? new FrameWriterParquetParallel() : new FrameWriterParquet(); case DELTA: return new FrameWriterDelta(); default: diff --git a/src/main/java/org/apache/sysds/runtime/io/FrameWriterParquet.java b/src/main/java/org/apache/sysds/runtime/io/FrameWriterParquet.java index ccaeeb56d51..775bfde0ce7 100644 --- a/src/main/java/org/apache/sysds/runtime/io/FrameWriterParquet.java +++ b/src/main/java/org/apache/sysds/runtime/io/FrameWriterParquet.java @@ -19,19 +19,23 @@ package org.apache.sysds.runtime.io; import java.io.IOException; -import java.util.ArrayList; -import java.util.List; +import java.util.HashMap; +import java.util.Map; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.FileSystem; import org.apache.hadoop.fs.Path; import org.apache.hadoop.mapred.JobConf; -import org.apache.parquet.example.data.Group; -import org.apache.parquet.example.data.simple.SimpleGroupFactory; +import org.apache.parquet.hadoop.ParquetOutputFormat; import org.apache.parquet.hadoop.ParquetWriter; -import org.apache.parquet.hadoop.example.ExampleParquetWriter; +import org.apache.parquet.hadoop.api.WriteSupport; +import org.apache.parquet.hadoop.metadata.CompressionCodecName; +import org.apache.parquet.io.api.Binary; +import org.apache.parquet.io.api.RecordConsumer; +import org.apache.parquet.schema.LogicalTypeAnnotation; import org.apache.parquet.schema.MessageType; -import org.apache.parquet.schema.MessageTypeParser; +import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName; +import org.apache.parquet.schema.Types; import org.apache.sysds.conf.ConfigurationManager; import org.apache.sysds.runtime.DMLRuntimeException; import org.apache.sysds.runtime.frame.data.FrameBlock; @@ -44,6 +48,28 @@ */ public class FrameWriterParquet extends FrameWriter { + public enum DictEncoding { + ALL_ON, ALL_OFF, STRING_ONLY + } + + private final CompressionCodecName codec; + private final DictEncoding dictEncoding; + private final long rowGroupSize; + + public FrameWriterParquet() { + this(CompressionCodecName.ZSTD, DictEncoding.STRING_ONLY, ParquetWriter.DEFAULT_BLOCK_SIZE); + } + + public FrameWriterParquet(CompressionCodecName codec, DictEncoding dictEncoding) { + this(codec, dictEncoding, ParquetWriter.DEFAULT_BLOCK_SIZE); + } + + public FrameWriterParquet(CompressionCodecName codec, DictEncoding dictEncoding, long rowGroupSize) { + this.codec = codec; + this.dictEncoding = dictEncoding; + this.rowGroupSize = rowGroupSize; + } + /** * Writes a FrameBlock to a Parquet file on HDFS. * @@ -71,9 +97,9 @@ public final void writeFrameToHDFS(FrameBlock src, String fname, long rlen, long } /** - * Writes the FrameBlock data to a Parquet file using a ParquetWriter. - * The method generates a Parquet schema based on the metadata of the FrameBlock, initializes a ParquetWriter with specified configurations, - * iterates over each row and column, adding values (in batches for improved performance) using type-specific conversions. + * Writes the FrameBlock data to a Parquet file using a ParquetWriter. The method generates a Parquet schema based + * on the metadata of the FrameBlock, initializes a ParquetWriter with specified configurations, iterates over each + * row and column, writing directly to the RecordConsumer, using type-specific conversions. * * @param path The HDFS path where the Parquet file will be written. * @param conf The Hadoop configuration. @@ -87,70 +113,27 @@ protected void writeParquetFrameToHDFS(Path path, Configuration conf, FrameBlock // Create schema based on frame block metadata MessageType schema = createParquetSchema(src); - // TODO:Experiment with different batch sizes? - int batchSize = 1000; - int rowCount = 0; - - // Write data using ParquetWriter //FIXME replace example writer? - try (ParquetWriter writer = ExampleParquetWriter.builder(path) - .withConf(conf) - .withType(schema) - .withCompressionCodec(ParquetWriter.DEFAULT_COMPRESSION_CODEC_NAME) - .withRowGroupSize((long) ParquetWriter.DEFAULT_BLOCK_SIZE) - .withPageSize(ParquetWriter.DEFAULT_PAGE_SIZE) - .withDictionaryEncoding(true) - .build()) - { - - SimpleGroupFactory groupFactory = new SimpleGroupFactory(schema); - - List rowBuffer = new ArrayList<>(batchSize); - - for (int i = 0; i < src.getNumRows(); i++) { - Group group = groupFactory.newGroup(); - for (int j = 0; j < src.getNumColumns(); j++) { - Object value = src.get(i, j); - if (value != null) { - ValueType type = src.getSchema()[j]; - switch (type) { - case STRING: - group.add(src.getColumnNames()[j], value.toString()); - break; - case INT32: - group.add(src.getColumnNames()[j], (int) value); - break; - case INT64: - group.add(src.getColumnNames()[j], (long) value); - break; - case FP32: - group.add(src.getColumnNames()[j], (float) value); - break; - case FP64: - group.add(src.getColumnNames()[j], (double) value); - break; - case BOOLEAN: - group.add(src.getColumnNames()[j], (boolean) value); - break; - default: - throw new IOException("Unsupported value type: " + type); - } - } - } - rowBuffer.add(group); - rowCount++; + String[] columnNames = src.getColumnNames(); + ValueType[] columnTypes = src.getSchema(); - if (rowCount >= batchSize) { - for (Group g : rowBuffer) { - writer.write(g); - } - rowBuffer.clear(); - rowCount = 0; - } - } - - for (Group g : rowBuffer) { - writer.write(g); - } + FrameParquetWriterBuilder writerBuilder = new FrameParquetWriterBuilder(path, schema, src).withConf(conf) + .withCompressionCodec( + CompressionCodecName.fromConf(conf.get(ParquetOutputFormat.COMPRESSION, codec.name()))) + .withRowGroupSize(conf.getLong(ParquetOutputFormat.BLOCK_SIZE, rowGroupSize)) + .withPageSize(conf.getInt(ParquetOutputFormat.PAGE_SIZE, ParquetWriter.DEFAULT_PAGE_SIZE)) + .withDictionaryPageSize( + conf.getInt(ParquetOutputFormat.DICTIONARY_PAGE_SIZE, ParquetWriter.DEFAULT_PAGE_SIZE)) + .withDictionaryEncoding( + conf.getBoolean(ParquetOutputFormat.ENABLE_DICTIONARY, dictEncoding == DictEncoding.ALL_ON)); + + if(dictEncoding == DictEncoding.STRING_ONLY) + for(int j = 0; j < src.getNumColumns(); j++) + if(columnTypes[j] == ValueType.STRING) + writerBuilder = writerBuilder.withDictionaryEncoding(columnNames[j], true); + + try(ParquetWriter writer = writerBuilder.build()) { + for(int i = 0; i < src.getNumRows(); i++) + writer.write(i); } // Delete CRC files created by Hadoop if necessary @@ -164,36 +147,126 @@ protected void writeParquetFrameToHDFS(Path path, Configuration conf, FrameBlock * @return The generated Parquet MessageType schema. */ protected MessageType createParquetSchema(FrameBlock src) { - StringBuilder schemaBuilder = new StringBuilder("message FrameSchema {"); String[] columnNames = src.getColumnNames(); ValueType[] columnTypes = src.getSchema(); + Types.MessageTypeBuilder builder = Types.buildMessage(); for (int i = 0; i < src.getNumColumns(); i++) { - schemaBuilder.append("optional "); switch (columnTypes[i]) { case STRING: - schemaBuilder.append("binary ").append(columnNames[i]).append(" (UTF8);"); + builder.optional(PrimitiveTypeName.BINARY).as(LogicalTypeAnnotation.stringType()) + .named(columnNames[i]); break; case INT32: - schemaBuilder.append("int32 ").append(columnNames[i]).append(";"); + builder.optional(PrimitiveTypeName.INT32).named(columnNames[i]); break; case INT64: - schemaBuilder.append("int64 ").append(columnNames[i]).append(";"); + builder.optional(PrimitiveTypeName.INT64).named(columnNames[i]); break; case FP32: - schemaBuilder.append("float ").append(columnNames[i]).append(";"); + builder.optional(PrimitiveTypeName.FLOAT).named(columnNames[i]); break; case FP64: - schemaBuilder.append("double ").append(columnNames[i]).append(";"); + builder.optional(PrimitiveTypeName.DOUBLE).named(columnNames[i]); break; case BOOLEAN: - schemaBuilder.append("boolean ").append(columnNames[i]).append(";"); + builder.optional(PrimitiveTypeName.BOOLEAN).named(columnNames[i]); break; default: throw new IllegalArgumentException("Unsupported data type: " + columnTypes[i]); } } - schemaBuilder.append("}"); - return MessageTypeParser.parseMessageType(schemaBuilder.toString()); + return builder.named("FrameSchema"); + } + + /** + * WriteSupport implementation that writes rows from a FrameBlock directly to the Parquet RecordConsumer. + */ + private static class FrameWriteSupport extends WriteSupport { + private final MessageType schema; + private final FrameBlock src; + private RecordConsumer recordConsumer; + // constant across all rows + private String[] colNames; + private ValueType[] colTypes; + private int numCols; + + FrameWriteSupport(MessageType schema, FrameBlock src) { + this.schema = schema; + this.src = src; + } + + @Override + public WriteContext init(Configuration configuration) { + Map metadata = new HashMap<>(); + return new WriteContext(schema, metadata); + } + + @Override + public void prepareForWrite(RecordConsumer consumer) { + this.recordConsumer = consumer; + this.colNames = src.getColumnNames(); + this.colTypes = src.getSchema(); + this.numCols = src.getNumColumns(); + } + + @Override + public void write(Integer rowIndex) { + recordConsumer.startMessage(); + for(int j = 0; j < numCols; j++) { + Object value = src.get(rowIndex, j); + if(value != null) { + recordConsumer.startField(colNames[j], j); + switch(colTypes[j]) { + case STRING: + recordConsumer.addBinary(Binary.fromString(value.toString())); + break; + case INT32: + recordConsumer.addInteger((int) value); + break; + case INT64: + recordConsumer.addLong((long) value); + break; + case FP32: + recordConsumer.addFloat((float) value); + break; + case FP64: + recordConsumer.addDouble((double) value); + break; + case BOOLEAN: + recordConsumer.addBoolean((boolean) value); + break; + default: + throw new IllegalArgumentException("Unsupported value type: " + colTypes[j]); + } + recordConsumer.endField(colNames[j], j); + } + } + recordConsumer.endMessage(); + } + } + + /** + * ParquetWriter builder wired to FrameWriteSupport. + */ + private static class FrameParquetWriterBuilder extends ParquetWriter.Builder { + private final MessageType schema; + private final FrameBlock src; + + FrameParquetWriterBuilder(Path path, MessageType schema, FrameBlock src) { + super(path); + this.schema = schema; + this.src = src; + } + + @Override + protected FrameParquetWriterBuilder self() { + return this; + } + + @Override + protected WriteSupport getWriteSupport(Configuration conf) { + return new FrameWriteSupport(schema, src); + } } } diff --git a/src/main/java/org/apache/sysds/runtime/io/FrameWriterParquetParallel.java b/src/main/java/org/apache/sysds/runtime/io/FrameWriterParquetParallel.java index 0ef4431ef47..3efbbbefef0 100644 --- a/src/main/java/org/apache/sysds/runtime/io/FrameWriterParquetParallel.java +++ b/src/main/java/org/apache/sysds/runtime/io/FrameWriterParquetParallel.java @@ -35,17 +35,16 @@ import org.apache.sysds.utils.stats.InfrastructureAnalyzer; /** - * Multi-threaded frame parquet reader. + * Multi-threaded frame parquet writer. * */ public class FrameWriterParquetParallel extends FrameWriterParquet { /** - * Writes the FrameBlock data to HDFS in parallel. - * The method estimates the number of output partitions by comparing the total number of cells in the FrameBlock with the - * HDFS block size. It then determines the number of threads to use based on the parallelism configuration and the - * number of partitions. In case of parallelism, it divides the FrameBlock into chunks and a thread pool is created to - * execute a write task for each partition concurrently. + * Writes the FrameBlock data to HDFS in parallel. The method estimates the number of output partitions by comparing + * the estimated output size of the FrameBlock with the HDFS block size. It then determines the number of threads to + * use based on the parallelism configuration and the number of partitions. In case of parallelism, it divides the + * FrameBlock into chunks and a thread pool is created to execute a write task for each partition concurrently. * * @param path The HDFS path where the Parquet files will be written. * @param conf The Hadoop configuration. @@ -55,14 +54,15 @@ public class FrameWriterParquetParallel extends FrameWriterParquet { protected void writeParquetFrameToHDFS(Path path, Configuration conf, FrameBlock src) throws IOException, DMLRuntimeException { - // Estimate number of output partitions - int numPartFiles = Math.max((int) (src.getNumRows() * src.getNumColumns() / InfrastructureAnalyzer.getHDFSBlockSize()), 1); - + // Estimate output partitions from output size in bytes + int numPartFiles = Math + .max((int) (OptimizerUtils.estimateSizeExactFrame(src.getNumRows(), src.getNumColumns()) / + InfrastructureAnalyzer.getHDFSBlockSize()), 1); + // Determine parallelism int numThreads = Math.min(OptimizerUtils.getParallelBinaryWriteParallelism(), numPartFiles); - // Fall back to sequential write if numThreads <= 1 - if (numThreads <= 1) { + if(!_forcedParallel && numThreads <= 1) { super.writeParquetFrameToHDFS(path, conf, src); return; } diff --git a/src/test/java/org/apache/sysds/performance/frame/ParquetReaderBenchmark.java b/src/test/java/org/apache/sysds/performance/frame/ParquetReaderBenchmark.java new file mode 100644 index 00000000000..b2f4c77bcbd --- /dev/null +++ b/src/test/java/org/apache/sysds/performance/frame/ParquetReaderBenchmark.java @@ -0,0 +1,193 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.sysds.performance.frame; + +import java.io.BufferedReader; +import java.io.File; +import java.io.FileReader; +import java.io.FileWriter; +import java.io.PrintWriter; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +import org.apache.sysds.common.Types.ValueType; +import org.apache.sysds.runtime.frame.data.FrameBlock; +import org.apache.sysds.runtime.io.FrameReaderParquet; +import org.apache.sysds.runtime.io.FrameReaderParquetParallel; +import org.apache.sysds.runtime.io.FrameWriterParquet; +import org.junit.After; +import org.junit.Assume; + +/** + * Parquet reader benchmark comparing the sequential and parallel column-API readers on the TPC-H lineitem dataset. + * + * Results report median and min/max across runs and are appended to temp/benchmark_results.csv for plotting. + * + * The benchmark methods are disabled by default; uncomment the Test annotations to run manually. If the dataset is not + * present the benchmark is skipped with instructions. + */ +public class ParquetReaderBenchmark { + + private static final String TPCH_FILE = "temp/lineitem.tbl"; + private static final String RESULTS_CSV = "temp/benchmark_results.csv"; + private static final String TEMP_FILE = System.getProperty("java.io.tmpdir") + "/systemds_read_bench.parquet"; + private static final int RUNS = 7; + // override on larger machine with -DmaxRows=... (e.g. -DmaxRows=20000000) to test larger row groups. + private static final int MAX_ROWS = Integer.getInteger("maxRows", 2_000_000); + + // TPC-H lineitem schema + private static final ValueType[] LINEITEM_SCHEMA = {ValueType.INT64, ValueType.INT64, ValueType.INT64, + ValueType.INT32, ValueType.FP64, ValueType.FP64, ValueType.FP64, ValueType.FP64, ValueType.STRING, + ValueType.STRING, ValueType.STRING, ValueType.STRING, ValueType.STRING, ValueType.STRING, ValueType.STRING, + ValueType.STRING}; + private static final String[] LINEITEM_NAMES = {"orderkey", "partkey", "suppkey", "linenumber", "quantity", + "extendedprice", "discount", "tax", "returnflag", "linestatus", "shipdate", "commitdate", "receiptdate", + "shipinstruct", "shipmode", "comment"}; + + private PrintWriter csv; + + @After + public void cleanup() { + new File(TEMP_FILE).delete(); + if(csv != null) + csv.close(); + } + + // @Test + public void benchmarkReadTpch() throws Exception { + ReadSpec spec = writeTempParquet(loadLineitemOrSkip()); + runReadBenchmark("read_tpch", spec); + } + + private static final class ReadSpec { + final ValueType[] schema; + final String[] names; + final int rows; + final int cols; + + ReadSpec(ValueType[] schema, String[] names, int rows, int cols) { + this.schema = schema; + this.names = names; + this.rows = rows; + this.cols = cols; + } + } + + private ReadSpec writeTempParquet(FrameBlock data) throws Exception { + int rows = data.getNumRows(), cols = data.getNumColumns(); + ReadSpec spec = new ReadSpec(data.getSchema(), data.getColumnNames(), rows, cols); + new File(TEMP_FILE).delete(); + new FrameWriterParquet().writeFrameToHDFS(data, TEMP_FILE, rows, cols); + return spec; + } + + /** + * Writes a frame to Parquet once, then times reading it back with both readers. + */ + private void runReadBenchmark(String category, ReadSpec spec) throws Exception { + final ValueType[] schema = spec.schema; + final String[] names = spec.names; + final int rows = spec.rows, cols = spec.cols; + + FrameBlock probe = new FrameReaderParquet().readFrameFromHDFS(TEMP_FILE, schema, names, rows, cols); + org.junit.Assert.assertEquals("Row count mismatch", rows, probe.getNumRows()); + org.junit.Assert.assertEquals("Column count mismatch", cols, probe.getNumColumns()); + probe = null; + + openCsv(); + System.out.println( + "\n=== Parquet Read Benchmark [" + category + "] (" + rows + " rows, median of " + RUNS + " runs) ===\n"); + System.out.printf("%-24s %12s %15s %s%n", "Reader", "Median (ms)", "Rows/sec", "[runs]"); + System.out.println("-".repeat(72)); + + timeRead(category, "Sequential", + () -> new FrameReaderParquet().readFrameFromHDFS(TEMP_FILE, schema, names, rows, cols), rows); + timeRead(category, "Parallel", + () -> new FrameReaderParquetParallel().readFrameFromHDFS(TEMP_FILE, schema, names, rows, cols), rows); + System.out.println(); + } + + private interface ReadAction { + FrameBlock run() throws Exception; + } + + private void timeRead(String category, String label, ReadAction action, int rows) throws Exception { + action.run(); // warmup + + long[] times = new long[RUNS]; + for(int run = 0; run < RUNS; run++) { + long start = System.currentTimeMillis(); + action.run(); + times[run] = System.currentTimeMillis() - start; + } + long med = median(times); + long min = Arrays.stream(times).min().orElse(med); + long max = Arrays.stream(times).max().orElse(med); + System.out.printf("%-24s %12d %15.0f %14s %s%n", label, med, rows * 1000.0 / med, meanStd(times), + Arrays.toString(times)); + // columns: benchmark,label,time_ms(median),rows_per_sec,min_ms,max_ms + csv.printf("%s,%s,%d,%.0f,%d,%d%n", category, label, med, rows * 1000.0 / med, min, max); + } + + private static long median(long[] times) { + long[] sorted = times.clone(); + Arrays.sort(sorted); + return sorted[sorted.length / 2]; + } + + private static String meanStd(long[] times) { + double mean = Arrays.stream(times).average().orElse(0); + double var = Arrays.stream(times).mapToDouble(t -> (t - mean) * (t - mean)).average().orElse(0); + return String.format("%.0f+-%.0f ms", mean, Math.sqrt(var)); + } + + private void openCsv() throws Exception { + new File("temp").mkdirs(); + boolean exists = new File(RESULTS_CSV).exists(); + csv = new PrintWriter(new FileWriter(RESULTS_CSV, true)); + if(!exists) + csv.println("benchmark,label,time_ms,rows_per_sec,size_mb,compression_ratio"); + csv.flush(); + } + + private FrameBlock loadLineitemOrSkip() throws Exception { + File f = new File(TPCH_FILE); + if(!f.exists()) { + System.out.println("=== TPC-H read benchmark skipped, dataset not found at " + TPCH_FILE + " ==="); + Assume.assumeTrue("TPC-H dataset not found at " + TPCH_FILE, false); + } + System.out.print("Loading " + f.getPath() + " ... "); + List rows = new ArrayList<>(); + try(BufferedReader br = new BufferedReader(new FileReader(f))) { + String line; + while((line = br.readLine()) != null && rows.size() < MAX_ROWS) { + if(line.isEmpty()) + continue; + if(line.endsWith("|")) + line = line.substring(0, line.length() - 1); + rows.add(line.split("\\|", -1)); + } + } + String[][] arr = rows.toArray(new String[0][]); + System.out.println(arr.length + " rows loaded."); + return new FrameBlock(LINEITEM_SCHEMA, LINEITEM_NAMES, arr); + } + +} diff --git a/src/test/java/org/apache/sysds/performance/frame/ParquetWriterBenchmark.java b/src/test/java/org/apache/sysds/performance/frame/ParquetWriterBenchmark.java new file mode 100644 index 00000000000..aec1f05f70d --- /dev/null +++ b/src/test/java/org/apache/sysds/performance/frame/ParquetWriterBenchmark.java @@ -0,0 +1,206 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.sysds.performance.frame; + +import java.io.BufferedReader; +import java.io.File; +import java.io.FileReader; +import java.io.FileWriter; +import java.io.PrintWriter; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +import org.apache.parquet.hadoop.metadata.CompressionCodecName; +import org.apache.sysds.common.Types.ValueType; +import org.apache.sysds.runtime.frame.data.FrameBlock; +import org.apache.sysds.runtime.io.FrameWriterParquet; +import org.apache.sysds.runtime.io.FrameWriterParquet.DictEncoding; +import org.junit.After; +import org.junit.Assume; + +/** + * Parquet writer benchmark using the TPC-H lineitem dataset. Writes results to temp/benchmark_results.csv for plotting. + * + * The benchmark methods are disabled by default; uncomment to run manually. If the dataset is not present inside the + * temp dir the corresponding benchmark is skipped with instructions. + * + * The default maxRows=2_000_000, to benchmark row-group size properly, run on a machine with enough RAM for a larger + * frame: # 1. Generate a larger TPC-H lineitem into temp/lineitem.tbl # 2. Raise the maxRows cap, then run the + * benchmark: mvn test -Dtest=ParquetWriterBenchmark#benchmarkRowGroupSizes \ -DmaxRows=60000000 -DargLine="-Xms24g + * -Xmx24g" -DfailIfNoTests=false + */ +public class ParquetWriterBenchmark { + + private static final String TPCH_FILE = "temp/lineitem.tbl"; + private static final String RESULTS_CSV = "temp/benchmark_results.csv"; + private static final String TEMP_FILE = System.getProperty("java.io.tmpdir") + "/systemds_tpch_bench.parquet"; + private static final int RUNS = 3; + private static final int MAX_ROWS = Integer.getInteger("maxRows", 2_000_000); + private static final long[] ROW_GROUP_SIZES = {1024 * 1024, // 1 MB + 8L * 1024 * 1024, // 8 MB + 16L * 1024 * 1024, // 16 MB + 32L * 1024 * 1024, // 32 MB + 64L * 1024 * 1024, // 64 MB + 128L * 1024 * 1024, // 128 MB (Parquet default) + 256L * 1024 * 1024, // 256 MB + 512L * 1024 * 1024 // 512 MB + }; + + // TPC-H lineitem schema + private static final ValueType[] LINEITEM_SCHEMA = {ValueType.INT64, // orderkey + ValueType.INT64, // partkey + ValueType.INT64, // suppkey + ValueType.INT32, // linenumber + ValueType.FP64, // quantity + ValueType.FP64, // extendedprice + ValueType.FP64, // discount + ValueType.FP64, // tax + ValueType.STRING, // returnflag (3 unique values) + ValueType.STRING, // linestatus (2 unique values) + ValueType.STRING, // shipdate + ValueType.STRING, // commitdate + ValueType.STRING, // receiptdate + ValueType.STRING, // shipinstruct (4 unique values) + ValueType.STRING, // shipmode (7 unique values) + ValueType.STRING // comment + }; + + private static final String[] LINEITEM_NAMES = {"orderkey", "partkey", "suppkey", "linenumber", "quantity", + "extendedprice", "discount", "tax", "returnflag", "linestatus", "shipdate", "commitdate", "receiptdate", + "shipinstruct", "shipmode", "comment"}; + + private PrintWriter csv; + + @After + public void cleanup() { + new File(TEMP_FILE).delete(); + if(csv != null) + csv.close(); + } + + // @Test + public void benchmarkDictionaryEncoding() throws Exception { + FrameBlock data = loadOrSkip(); + int rows = data.getNumRows(); + + openCsv(); + System.out.println( + "\n=== TPC-H Dictionary Encoding Benchmark (" + rows + " rows, median of " + RUNS + " runs) ===\n"); + System.out.printf("%-20s %12s %15s%n", "Strategy", "Time (ms)", "Rows/sec"); + System.out.println("-".repeat(52)); + + time("encoding", "ALL_ON", new FrameWriterParquet(CompressionCodecName.UNCOMPRESSED, DictEncoding.ALL_ON), data, + rows); + time("encoding", "ALL_OFF", new FrameWriterParquet(CompressionCodecName.UNCOMPRESSED, DictEncoding.ALL_OFF), + data, rows); + time("encoding", "STRING_ONLY", + new FrameWriterParquet(CompressionCodecName.UNCOMPRESSED, DictEncoding.STRING_ONLY), data, rows); + System.out.println(); + } + + // @Test + public void benchmarkRowGroupSizes() throws Exception { + FrameBlock data = loadOrSkip(); + int rows = data.getNumRows(); + + openCsv(); + System.out + .println("\n=== TPC-H Row Group Size Benchmark (" + rows + " rows, median of " + RUNS + " runs) ===\n"); + System.out.printf("%-20s %12s %15s%n", "Row Group Size", "Time (ms)", "Rows/sec"); + System.out.println("-".repeat(52)); + + for(long rowGroupSize : ROW_GROUP_SIZES) { + String label = (rowGroupSize / (1024 * 1024)) + "MB"; + time("row_group_sizes", label, + new FrameWriterParquet(CompressionCodecName.ZSTD, DictEncoding.ALL_ON, rowGroupSize), data, rows); + } + System.out.println(); + } + + private FrameBlock loadOrSkip() throws Exception { + File f = new File(TPCH_FILE); + if(!f.exists()) { + System.out.println(); + System.out.println("==================================================="); + System.out.println("TPC-H benchmark skipped, dataset not found"); + System.out.println("To reproduce:"); + System.out.println(" 1. Install DuckDB: https://duckdb.org"); + System.out.println(" 2. Open shell: duckdb"); + System.out.println(" 3. Run in DuckDB: INSTALL tpch;"); + System.out.println(" LOAD tpch;"); + System.out.println(" CALL dbgen(sf=1);"); + System.out.println(" COPY lineitem TO '/temp/lineitem.tbl'"); + System.out.println(" (DELIMITER '|', HEADER false);"); + System.out.println("==================================================="); + Assume.assumeTrue("TPC-H dataset not found at " + TPCH_FILE, false); + } + return loadLineitem(f); + } + + private void openCsv() throws Exception { + new File("temp").mkdirs(); + boolean exists = new File(RESULTS_CSV).exists(); + csv = new PrintWriter(new FileWriter(RESULTS_CSV, true)); + if(!exists) + csv.println("benchmark,label,time_ms,rows_per_sec,size_mb,compression_ratio"); + csv.flush(); + } + + private void time(String category, String label, FrameWriterParquet writer, FrameBlock data, int rows) + throws Exception { + new File(TEMP_FILE).delete(); + writer.writeFrameToHDFS(data, TEMP_FILE, rows, data.getNumColumns()); // warmup + + long[] times = new long[RUNS]; + for(int run = 0; run < RUNS; run++) { + new File(TEMP_FILE).delete(); + long start = System.currentTimeMillis(); + writer.writeFrameToHDFS(data, TEMP_FILE, rows, data.getNumColumns()); + times[run] = System.currentTimeMillis() - start; + } + long med = median(times); + System.out.printf("%-20s %12d %15.0f%n", label, med, rows * 1000.0 / med); + csv.printf("%s,%s,%d,%.0f,,%n", category, label, med, rows * 1000.0 / med); + } + + private static long median(long[] times) { + long[] sorted = times.clone(); + Arrays.sort(sorted); + return sorted[sorted.length / 2]; + } + + private static FrameBlock loadLineitem(File f) throws Exception { + System.out.print("Loading " + f.getPath() + " ... "); + List rows = new ArrayList<>(); + try(BufferedReader br = new BufferedReader(new FileReader(f))) { + String line; + while((line = br.readLine()) != null && rows.size() < MAX_ROWS) { + if(line.isEmpty()) + continue; + if(line.endsWith("|")) + line = line.substring(0, line.length() - 1); + rows.add(line.split("\\|", -1)); + } + } + String[][] data = rows.toArray(new String[0][]); + System.out.println(data.length + " rows loaded."); + return new FrameBlock(LINEITEM_SCHEMA, LINEITEM_NAMES, data); + } +} diff --git a/src/test/java/org/apache/sysds/test/functions/io/parquet/FrameParquetSchemaTest.java b/src/test/java/org/apache/sysds/test/functions/io/parquet/FrameParquetSchemaTest.java index 1e4334891ed..cc1412b1606 100644 --- a/src/test/java/org/apache/sysds/test/functions/io/parquet/FrameParquetSchemaTest.java +++ b/src/test/java/org/apache/sysds/test/functions/io/parquet/FrameParquetSchemaTest.java @@ -19,9 +19,9 @@ package org.apache.sysds.test.functions.io.parquet; -import org.junit.Assert; import org.junit.Test; +import org.apache.sysds.test.TestUtils; import org.apache.sysds.common.Types.ValueType; import org.apache.sysds.runtime.frame.data.FrameBlock; import org.apache.sysds.runtime.io.FrameReader; @@ -55,169 +55,48 @@ public void setUp() { /** * Test for sequential writer and reader - * + * */ @Test - public void testParquetWriteReadAllSchemaTypes() { - String fname = output("Rout"); - - // Define a schema with one column per type - ValueType[] schema = new ValueType[] { - ValueType.FP64, - ValueType.FP32, - ValueType.INT32, - ValueType.INT64, - ValueType.BOOLEAN, - ValueType.STRING - }; - - // Create an empty frame block with the above schema - FrameBlock fb = new FrameBlock(schema); - - // Populate frame block - Object[][] rows = new Object[][] { - { 1.0, 1.1f, 10, 100L, true, "A" }, - { 2.0, 2.1f, 20, 200L, false, "B" }, - { 3.0, 3.1f, 30, 300L, true, "C" }, - { 4.0, 4.1f, 40, 400L, false, "D" }, - { 5.0, 5.1f, 50, 500L, true, "E" } - }; - - for (Object[] row : rows) { - fb.appendRow(row); - } - - System.out.println(fb); - - int numRows = fb.getNumRows(); - int numCols = fb.getNumColumns(); - - // Write the FrameBlock to a Parquet file using the sequential writer - try { - FrameWriter writer = new FrameWriterParquet(); - writer.writeFrameToHDFS(fb, fname, numRows, numCols); - } - catch (IOException e) { - e.printStackTrace(); - Assert.fail("Failed to write frame block to Parquet: " + e.getMessage()); - } - - // Read the Parquet file back into a new FrameBlock - FrameBlock fbRead = null; - try { - FrameReader reader = new FrameReaderParquet(); - String[] colNames = fb.getColumnNames(); - fbRead = reader.readFrameFromHDFS(fname, schema, colNames, numRows, numCols); - } - catch (IOException e) { - e.printStackTrace(); - Assert.fail("Failed to read frame block from Parquet: " + e.getMessage()); - } - - // Compare the original and the read frame blocks - compareFrameBlocks(fb, fbRead, 1e-6); + public void testParquetWriteReadAllSchemaTypes() throws IOException { + runWriteReadAllSchemaTypes(new FrameWriterParquet(), new FrameReaderParquet(), output("Rout")); } /** * Test for multithreaded writer and reader - * + * */ @Test - public void testParquetWriteReadAllSchemaTypesParallel() { - String fname = output("Rout_parallel"); - - ValueType[] schema = new ValueType[] { - ValueType.FP64, - ValueType.FP32, - ValueType.INT32, - ValueType.INT64, - ValueType.BOOLEAN, - ValueType.STRING - }; + public void testParquetWriteReadAllSchemaTypesParallel() throws IOException { + runWriteReadAllSchemaTypes(new FrameWriterParquetParallel(), new FrameReaderParquetParallel(), + output("Rout_parallel")); + } + + private static void runWriteReadAllSchemaTypes(FrameWriter writer, FrameReader reader, String fname) + throws IOException { + // Define a schema with one column per type + ValueType[] schema = new ValueType[] {ValueType.FP64, ValueType.FP32, ValueType.INT32, ValueType.INT64, + ValueType.BOOLEAN, ValueType.STRING}; + // Create an empty frame block with the above schema FrameBlock fb = new FrameBlock(schema); - Object[][] rows = new Object[][] { - { 1.0, 1.1f, 10, 100L, true, "A" }, - { 2.0, 2.1f, 20, 200L, false, "B" }, - { 3.0, 3.1f, 30, 300L, true, "C" }, - { 4.0, 4.1f, 40, 400L, false, "D" }, - { 5.0, 5.1f, 50, 500L, true, "E" } - }; + // Populate frame block + Object[][] rows = new Object[][] {{1.0, 1.1f, 10, 100L, true, "A"}, {2.0, 2.1f, 20, 200L, false, "B"}, + {3.0, 3.1f, 30, 300L, true, "C"}, {4.0, 4.1f, 40, 400L, false, "D"}, {5.0, 5.1f, 50, 500L, true, "E"}}; - for (Object[] row : rows) { + for(Object[] row : rows) fb.appendRow(row); - } int numRows = fb.getNumRows(); int numCols = fb.getNumColumns(); - try { - FrameWriter writer = new FrameWriterParquetParallel(); - writer.writeFrameToHDFS(fb, fname, numRows, numCols); - } - catch (IOException e) { - e.printStackTrace(); - Assert.fail("Failed to write frame block to Parquet (parallel): " + e.getMessage()); - } - - FrameBlock fbRead = null; - try { - FrameReader reader = new FrameReaderParquetParallel(); - String[] colNames = fb.getColumnNames(); - fbRead = reader.readFrameFromHDFS(fname, schema, colNames, numRows, numCols); - } - catch (IOException e) { - e.printStackTrace(); - Assert.fail("Failed to read frame block from Parquet (parallel): " + e.getMessage()); - } - - compareFrameBlocks(fb, fbRead, 1e-6); - } + writer.writeFrameToHDFS(fb, fname, numRows, numCols); - private void compareFrameBlocks(FrameBlock expected, FrameBlock actual, double eps) { - Assert.assertEquals("Number of rows mismatch", expected.getNumRows(), actual.getNumRows()); - Assert.assertEquals("Number of columns mismatch", expected.getNumColumns(), actual.getNumColumns()); - - int rows = expected.getNumRows(); - int cols = expected.getNumColumns(); - - for (int i = 0; i < rows; i++) { - for (int j = 0; j < cols; j++) { - Object expVal = expected.get(i, j); - Object actVal = actual.get(i, j); - ValueType vt = expected.getSchema()[j]; - - // Handle nulls first - if(expVal == null || actVal == null) { - Assert.assertEquals("Mismatch at (" + i + "," + j + ")", expVal, actVal); - } else { - switch(vt) { - case FP64: - case FP32: - double dExp = ((Number) expVal).doubleValue(); - double dAct = ((Number) actVal).doubleValue(); - Assert.assertEquals("Mismatch at (" + i + "," + j + ")", dExp, dAct, eps); - break; - case INT32: - case INT64: - long lExp = ((Number) expVal).longValue(); - long lAct = ((Number) actVal).longValue(); - Assert.assertEquals("Mismatch at (" + i + "," + j + ")", lExp, lAct); - break; - case BOOLEAN: - boolean bExp = (Boolean) expVal; - boolean bAct = (Boolean) actVal; - Assert.assertEquals("Mismatch at (" + i + "," + j + ")", bExp, bAct); - break; - case STRING: - Assert.assertEquals("Mismatch at (" + i + "," + j + ")", expVal.toString(), actVal.toString()); - break; - default: - Assert.fail("Unsupported type in comparison: " + vt); - } - } - } - } + String[] colNames = fb.getColumnNames(); + FrameBlock fbRead = reader.readFrameFromHDFS(fname, schema, colNames, numRows, numCols); + + // Compare the original and the read frame blocks + TestUtils.compareFrames(fb, fbRead, false); } } diff --git a/src/test/java/org/apache/sysds/test/functions/io/parquet/FrameReaderWriterParquetTest.java b/src/test/java/org/apache/sysds/test/functions/io/parquet/FrameReaderWriterParquetTest.java new file mode 100644 index 00000000000..0a8f9156122 --- /dev/null +++ b/src/test/java/org/apache/sysds/test/functions/io/parquet/FrameReaderWriterParquetTest.java @@ -0,0 +1,152 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.test.functions.io.parquet; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.fail; + +import java.io.IOException; + +import org.apache.sysds.common.Types.ValueType; +import org.apache.sysds.runtime.frame.data.FrameBlock; +import org.apache.sysds.runtime.io.FrameReaderParquet; +import org.apache.sysds.runtime.io.FrameWriterParquet; +import org.apache.sysds.test.TestUtils; +import org.junit.Test; + +/** + * Random-frame write/read round-trip test for the parquet frame reader/writer. + */ +public class FrameReaderWriterParquetTest { + + private static final String FILENAME = "target/testTemp/functions/io/parquet/FrameReaderWriterParquetTest/frame.parquet"; + + // Parquet-supported value types + private static final ValueType[] SCHEMA = {ValueType.FP64, ValueType.FP32, ValueType.INT32, ValueType.INT64, + ValueType.BOOLEAN, ValueType.STRING}; + + @Test + public void testSingleRowSingleCol() throws IOException { + runWriteReadRoundTrip(1, 1, 4669201); + } + + @Test + public void testSingleRowMultiCol() throws IOException { + runWriteReadRoundTrip(1, 6, 4669201); + } + + @Test + public void testMultiRowSingleCol() throws IOException { + runWriteReadRoundTrip(21, 1, 4669201); + } + + @Test + public void testMultiRowMultiCol() throws IOException { + runWriteReadRoundTrip(42, 5, 4669201); + } + + @Test + public void testLargerFrame() throws IOException { + runWriteReadRoundTrip(694, 6, 4669201); + } + + @Test + public void testValueTypeEdgeCases() throws IOException { + // type min/max, empty string, special chars (comma/quote/newline/unicode), column name with space + ValueType[] schema = {ValueType.FP32, ValueType.FP64, ValueType.INT32, ValueType.INT64, ValueType.BOOLEAN, + ValueType.STRING}; + String[] names = {"f 32", "f64", "i32", "i64", "b", "s"}; + String[][] data = { + {String.valueOf(Float.MAX_VALUE), String.valueOf(Double.MAX_VALUE), String.valueOf(Integer.MAX_VALUE), + String.valueOf(Long.MAX_VALUE), "true", ""}, + {String.valueOf(-Float.MAX_VALUE), String.valueOf(-Double.MAX_VALUE), String.valueOf(Integer.MIN_VALUE), + String.valueOf(Long.MIN_VALUE), "false", "a,b\"c\nd"}, + {"0.0", "0.0", "0", "0", "true", "unicode_é中"}}; + FrameBlock in = new FrameBlock(schema, names, data); + + new FrameWriterParquet().writeFrameToHDFS(in, FILENAME, 3, 6); + FrameBlock out = new FrameReaderParquet().readFrameFromHDFS(FILENAME, schema, names, 3, 6); + + TestUtils.compareFrames(in, out, false); + } + + @Test + public void testNullsInStringColumn() throws IOException { + // Numeric columns from String[][] convert null to 0. Only test string nulls here. + // Numeric nulls are covered by the Spark-written userdata1/all files in ReadParquetTest. + ValueType[] schema = {ValueType.STRING, ValueType.STRING}; + String[] names = {"a", "b"}; + String[][] data = {{"x", "y"}, {null, null}, {"p", null}}; + FrameBlock in = new FrameBlock(schema, names, data); + + new FrameWriterParquet().writeFrameToHDFS(in, FILENAME, 3, 2); + FrameBlock out = new FrameReaderParquet().readFrameFromHDFS(FILENAME, schema, names, 3, 2); + + org.junit.Assert.assertNull(out.get(1, 0)); + org.junit.Assert.assertNull(out.get(1, 1)); + org.junit.Assert.assertNull(out.get(2, 1)); + org.junit.Assert.assertNotNull(out.get(0, 0)); + org.junit.Assert.assertNotNull(out.get(2, 0)); + TestUtils.compareFrames(in, out, false); + } + + @Test + public void testMismatchedSchemaRead() throws IOException { + // requested frame types differ from the parquet physical types -> per-cell conversion path + ValueType[] writeSchema = {ValueType.INT32, ValueType.FP32, ValueType.BOOLEAN, ValueType.FP64, + ValueType.STRING}; + String[] names = {"i32", "f32", "b", "f64", "s"}; + String[][] data = {{"1", "1.5", "true", "2.5", "7"}, {"2", "2.5", "false", "3.5", null}}; + FrameBlock in = new FrameBlock(writeSchema, names, data); + + new FrameWriterParquet().writeFrameToHDFS(in, FILENAME, 2, 5); + + ValueType[] readSchema = {ValueType.INT64, ValueType.FP64, ValueType.STRING, ValueType.STRING, ValueType.FP64}; + FrameBlock out = new FrameReaderParquet().readFrameFromHDFS(FILENAME, readSchema, names, 2, 5); + + assertEquals(1L, out.get(0, 0)); + assertEquals(1.5, ((Number) out.get(0, 1)).doubleValue(), 0.0); + assertEquals("true", out.get(0, 2)); + assertEquals("2.5", out.get(0, 3)); + assertEquals(7.0, ((Number) out.get(0, 4)).doubleValue(), 0.0); + // numeric nulls keep the array default + assertEquals(0.0, ((Number) out.get(1, 4)).doubleValue(), 0.0); + } + + private void runWriteReadRoundTrip(int rows, int cols, long seed) throws IOException { + try { + ValueType[] schema = new ValueType[cols]; + for(int i = 0; i < cols; i++) + schema[i] = SCHEMA[i % SCHEMA.length]; + + FrameBlock writeBlock = TestUtils.generateRandomFrameBlock(rows, schema, seed); + + new FrameWriterParquet().writeFrameToHDFS(writeBlock, FILENAME, rows, cols); + FrameBlock readBlock = new FrameReaderParquet().readFrameFromHDFS(FILENAME, schema, + writeBlock.getColumnNames(), rows, cols); + + TestUtils.compareFrames(writeBlock, readBlock, false); + } + catch(Exception e) { + e.printStackTrace(); + fail(e.getMessage()); + } + } +} diff --git a/src/test/java/org/apache/sysds/test/functions/io/parquet/ParquetTestUtils.java b/src/test/java/org/apache/sysds/test/functions/io/parquet/ParquetTestUtils.java new file mode 100644 index 00000000000..0ae2145839c --- /dev/null +++ b/src/test/java/org/apache/sysds/test/functions/io/parquet/ParquetTestUtils.java @@ -0,0 +1,275 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.test.functions.io.parquet; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.StandardCopyOption; +import java.sql.Timestamp; +import java.util.Arrays; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.Path; +import org.apache.parquet.hadoop.ParquetFileReader; +import org.apache.parquet.hadoop.metadata.BlockMetaData; +import org.apache.parquet.hadoop.metadata.ParquetMetadata; +import org.apache.parquet.hadoop.util.HadoopInputFile; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.PrimitiveType; +import org.apache.spark.sql.DataFrameWriter; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.SaveMode; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import org.apache.sysds.common.Types.ValueType; +import org.apache.sysds.conf.ConfigurationManager; +import org.apache.sysds.runtime.frame.data.FrameBlock; +import org.apache.sysds.test.AutomatedTestBase; + +class ParquetTestUtils { + + static class ParquetMetadataInfo { + String[] names; + ValueType[] schema; + long rlen; + long clen; + } + + static ParquetMetadataInfo inferMetadata(String fname) throws IOException { + Configuration conf = ConfigurationManager.getCachedJobConf(); + Path path = new Path(fname); + + ParquetMetadata metadata; + try(ParquetFileReader r = ParquetFileReader.open(HadoopInputFile.fromPath(path, conf))) { + metadata = r.getFooter(); + } + MessageType parquetSchema = metadata.getFileMetaData().getSchema(); + + int fieldCount = parquetSchema.getFieldCount(); + String[] names = new String[fieldCount]; + ValueType[] schema = new ValueType[fieldCount]; + + for(int i = 0; i < fieldCount; i++) { + names[i] = parquetSchema.getFieldName(i); + PrimitiveType.PrimitiveTypeName type = parquetSchema.getType(i).asPrimitiveType().getPrimitiveTypeName(); + switch(type) { + case INT32: + schema[i] = ValueType.INT32; + break; + case INT64: + schema[i] = ValueType.INT64; + break; + case FLOAT: + schema[i] = ValueType.FP32; + break; + case DOUBLE: + schema[i] = ValueType.FP64; + break; + case BOOLEAN: + schema[i] = ValueType.BOOLEAN; + break; + case BINARY: + schema[i] = ValueType.STRING; + break; + default: + throw new IOException("Unsupported parquet type: " + type + " in column " + names[i]); + } + } + + long rlen = 0; + for(BlockMetaData block : metadata.getBlocks()) + rlen += block.getRowCount(); + + ParquetMetadataInfo info = new ParquetMetadataInfo(); + info.names = names; + info.schema = schema; + info.rlen = rlen; + info.clen = fieldCount; + return info; + } + + static SparkSession sparkSession() { + return AutomatedTestBase.createSystemDSSparkSession("parquet-frame-tests", "local[1]"); + } + + /** + * Reads a parquet file (or directory of part files) through Spark, into a FrameBlock with the given column layout, + * projecting the columns in the requested order. + * + * @param fname parquet file or directory + * @param schema value types of the requested columns + * @param names names of the requested columns + * @return the frame as read by Spark + */ + static FrameBlock sparkReadAsFrame(String fname, ValueType[] schema, String[] names) { + Dataset df = sparkSession().read().parquet(fname).select(names[0], + Arrays.copyOfRange(names, 1, names.length)); + return toFrameBlock(df.collectAsList(), schema, names); + } + + /** Converts collected Spark rows into a FrameBlock */ + static FrameBlock toFrameBlock(List rows, ValueType[] schema, String[] names) { + FrameBlock fb = new FrameBlock(schema, names); + fb.ensureAllocatedColumns(rows.size()); + for(int r = 0; r < rows.size(); r++) + for(int c = 0; c < schema.length; c++) + fb.set(r, c, rows.get(r).get(c)); + return fb; + } + + /** + * Generates the public test files (userdata1, alltypes_plain, all) with Spark's DataFrameWriter, covering all + * supported physical types and null values. One file is gzip-compressed to also cover a non-default codec on read; + * the others use Spark's default (snappy). + * + * @param outDir directory the generated files are written into + * @return map from file name (e.g. "userdata1") to its generated file path + */ + static Map generatePublicTestFiles(File outDir) throws Exception { + SparkSession spark = sparkSession(); + + Map files = new LinkedHashMap<>(); + files.put("userdata1", writeTestFile(spark, outDir, "userdata1", userdata1Rows(), userdata1Schema(), null)); + files.put("alltypes_plain", + writeTestFile(spark, outDir, "alltypes_plain", alltypesPlainRows(), alltypesPlainSchema(), "gzip")); + files.put("all", writeTestFile(spark, outDir, "all", allRows(), allSchema(), null)); + return files; + } + + /** + * Writes a file containing a deprecated INT96 timestamp column (id, ts, name), forced via + * spark.sql.parquet.outputTimestampType so the encoding does not depend on the Spark version. + * + * @param outDir directory the file is written into + * @return the generated file path + */ + static String writeInt96TimestampFile(File outDir) throws IOException { + SparkSession spark = sparkSession(); + StructType schema = new StructType( + new StructField[] {new StructField("id", DataTypes.IntegerType, true, Metadata.empty()), + new StructField("ts", DataTypes.TimestampType, true, Metadata.empty()), + new StructField("name", DataTypes.StringType, true, Metadata.empty()),}); + List rows = Arrays.asList(RowFactory.create(1, Timestamp.valueOf("2016-02-03 07:55:29"), "Amanda"), + RowFactory.create(2, Timestamp.valueOf("2016-02-03 17:04:03"), "Albert"), + RowFactory.create(3, null, "Evelyn")); + spark.conf().set("spark.sql.parquet.outputTimestampType", "INT96"); + try { + return writeTestFile(spark, outDir, "int96", rows, schema, null); + } + finally { + spark.conf().unset("spark.sql.parquet.outputTimestampType"); + } + } + + private static StructType userdata1Schema() { + return new StructType(new StructField[] {new StructField("id", DataTypes.IntegerType, true, Metadata.empty()), + new StructField("first_name", DataTypes.StringType, true, Metadata.empty()), + new StructField("salary", DataTypes.DoubleType, true, Metadata.empty()),}); + } + + private static List userdata1Rows() { + return Arrays.asList(RowFactory.create(1, "Amanda", 49756.53), RowFactory.create(2, "Albert", 150280.17), + RowFactory.create(3, "Evelyn", 144972.51), RowFactory.create(4, "Denise", null), + RowFactory.create(5, "Carlos", 75500.34)); + } + + private static StructType alltypesPlainSchema() { + return new StructType(new StructField[] {new StructField("id", DataTypes.IntegerType, true, Metadata.empty()), + new StructField("bool_col", DataTypes.BooleanType, true, Metadata.empty()), + new StructField("tinyint_col", DataTypes.IntegerType, true, Metadata.empty()), + new StructField("smallint_col", DataTypes.IntegerType, true, Metadata.empty()), + new StructField("bigint_col", DataTypes.LongType, true, Metadata.empty()), + new StructField("float_col", DataTypes.FloatType, true, Metadata.empty()), + new StructField("double_col", DataTypes.DoubleType, true, Metadata.empty()), + new StructField("date_string_col", DataTypes.StringType, true, Metadata.empty()), + new StructField("string_col", DataTypes.StringType, true, Metadata.empty()),}); + } + + private static List alltypesPlainRows() { + return Arrays.asList(RowFactory.create(1, true, 1, 10, 100L, 1.5f, 2.25, "03/01/09", "row-1"), + RowFactory.create(2, false, 2, 20, 200L, 2.5f, 4.5, "03/02/09", "row-2"), + RowFactory.create(3, true, 3, 30, 300L, 3.5f, 6.75, "03/03/09", "row-3"), + RowFactory.create(4, false, 4, 40, 400L, 4.5f, 9.0, "03/04/09", "row-4"), + RowFactory.create(5, true, 5, 50, 500L, 5.5f, 11.25, "03/05/09", "row-5"), + RowFactory.create(6, false, 6, 60, 600L, 6.5f, 13.5, "03/06/09", "row-6"), + RowFactory.create(7, true, 7, 70, 700L, 7.5f, 15.75, "03/07/09", "row-7"), + RowFactory.create(8, false, 8, 80, 800L, 8.5f, 18.0, "03/08/09", "row-8")); + } + + private static StructType allSchema() { + return new StructType( + new StructField[] {new StructField("PassengerId", DataTypes.IntegerType, true, Metadata.empty()), + new StructField("Survived", DataTypes.IntegerType, true, Metadata.empty()), + new StructField("Pclass", DataTypes.IntegerType, true, Metadata.empty()), + new StructField("Name", DataTypes.StringType, true, Metadata.empty()), + new StructField("Sex", DataTypes.StringType, true, Metadata.empty()), + new StructField("Age", DataTypes.DoubleType, true, Metadata.empty()), + new StructField("Fare", DataTypes.DoubleType, true, Metadata.empty()), + new StructField("Embarked", DataTypes.StringType, true, Metadata.empty()),}); + } + + private static List allRows() { + return Arrays.asList(RowFactory.create(1, 0, 3, "Braund, Mr. Owen Harris", "male", 22.0, 7.25, "S"), + RowFactory.create(2, 1, 1, "Cumings, Mrs. John Bradley", "female", 38.0, 71.2833, "C"), + RowFactory.create(3, 1, 3, "Heikkinen, Miss. Laina", "female", 26.0, 7.925, "S"), + RowFactory.create(4, 1, 1, "Futrelle, Mrs. Jacques Heath", "female", 35.0, 53.1, "S"), + RowFactory.create(5, 0, 3, "Allen, Mr. William Henry", "male", null, 8.05, "S"), + RowFactory.create(6, 0, 3, "Moran, Mr. James", "male", null, 8.4583, "Q"), + RowFactory.create(7, 0, 1, "McCarthy, Mr. Timothy J", "male", 54.0, 51.8625, "S"), + RowFactory.create(8, 0, 3, "Palsson, Master. Gosta Leonard", "male", 2.0, 21.075, "S")); + } + + // Spark writes a directory of part files, so we force one partition and rename it to a single file. + private static String writeTestFile(SparkSession spark, File outDir, String name, List rows, StructType schema, + String codec) throws IOException { + Dataset df = spark.createDataFrame(rows, schema); + File tmpDir = new File(outDir, "_tmp_" + name); + DataFrameWriter writer = df.coalesce(1).write().mode(SaveMode.Overwrite); + if(codec != null) + writer = writer.option("compression", codec); + writer.parquet(tmpDir.getPath()); + + File[] parts = tmpDir.listFiles((d, n) -> n.startsWith("part-") && n.endsWith(".parquet")); + if(parts == null || parts.length != 1) + throw new IOException("expected exactly 1 part file in " + tmpDir); + + File dest = new File(outDir, name + ".parquet"); + Files.copy(parts[0].toPath(), dest.toPath(), StandardCopyOption.REPLACE_EXISTING); + deleteRecursive(tmpDir); + return dest.getPath(); + } + + private static void deleteRecursive(File f) { + File[] children = f.listFiles(); + if(children != null) + for(File c : children) + deleteRecursive(c); + f.delete(); + } +} diff --git a/src/test/java/org/apache/sysds/test/functions/io/parquet/ReadParquetTest.java b/src/test/java/org/apache/sysds/test/functions/io/parquet/ReadParquetTest.java new file mode 100644 index 00000000000..3e5fc17c907 --- /dev/null +++ b/src/test/java/org/apache/sysds/test/functions/io/parquet/ReadParquetTest.java @@ -0,0 +1,188 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.test.functions.io.parquet; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Files; +import java.util.Map; + +import org.apache.spark.sql.SaveMode; +import org.apache.spark.sql.SparkSession; +import org.apache.sysds.common.Types.ValueType; +import org.apache.sysds.runtime.frame.data.FrameBlock; +import org.apache.sysds.runtime.io.FrameReaderParquet; +import org.apache.sysds.runtime.io.FrameReaderParquetParallel; +import org.apache.sysds.runtime.io.FrameWriterParquet; +import org.apache.sysds.test.functions.io.parquet.ParquetTestUtils.ParquetMetadataInfo; +import org.apache.sysds.test.TestUtils; +import org.junit.AfterClass; +import org.junit.Assert; +import org.junit.BeforeClass; +import org.junit.Test; + +/** + * Verifies the parquet frame readers against Spark's parquet reader on Spark-written files. + */ +public class ReadParquetTest { + + // Generated once per test class with Spark's DataFrameWriter + private static File testFileDir; + private static String[] FILENAMES; + + @BeforeClass + public static void generateTestFiles() throws Exception { + testFileDir = Files.createTempDirectory("systemds_parquet_public_test_files").toFile(); + Map files = ParquetTestUtils.generatePublicTestFiles(testFileDir); + FILENAMES = new String[] {files.get("userdata1"), files.get("alltypes_plain"), files.get("all")}; + } + + @AfterClass + public static void cleanupTestFiles() { + deleteRecursive(testFileDir); + } + + private static void deleteRecursive(File f) { + File[] children = f.listFiles(); + if(children != null) + for(File c : children) + deleteRecursive(c); + f.delete(); + } + + @Test + public void testReadMatchesSpark() throws Exception { + for(String filename : FILENAMES) { + ParquetMetadataInfo info = ParquetTestUtils.inferMetadata(filename); + + FrameBlock expected = ParquetTestUtils.sparkReadAsFrame(filename, info.schema, info.names); + FrameBlock actual = new FrameReaderParquet().readFrameFromHDFS(filename, info.schema, info.names, info.rlen, + info.clen); + + TestUtils.compareFrames(expected, actual, false); + } + } + + @Test + public void testColumnSubsetProjection() throws Exception { + // request a reordered subset of columns (d, a, c; skip b) + File temp = Files.createTempFile("systemds_subset_parquet", ".parquet").toFile(); + try { + ValueType[] fullSchema = {ValueType.INT64, ValueType.STRING, ValueType.FP64, ValueType.BOOLEAN}; + String[] fullNames = {"a", "b", "c", "d"}; + FrameBlock original = new FrameBlock(fullSchema, fullNames, + new String[][] {{"10", "x", "1.5", "true"}, {"20", "y", "2.5", "false"}, {"30", "z", "3.5", "true"}}); + new FrameWriterParquet().writeFrameToHDFS(original, temp.getPath(), 3, 4); + + ValueType[] subSchema = {ValueType.BOOLEAN, ValueType.INT64, ValueType.FP64}; + String[] subNames = {"d", "a", "c"}; + + FrameBlock expected = ParquetTestUtils.sparkReadAsFrame(temp.getPath(), subSchema, subNames); + FrameBlock actual = new FrameReaderParquet().readFrameFromHDFS(temp.getPath(), subSchema, subNames, 3, 3); + + TestUtils.compareFrames(expected, actual, false); + } + finally { + temp.delete(); + } + } + + @Test + public void testInt96ColumnRejectedOthersReadable() throws Exception { + String file = ParquetTestUtils.writeInt96TimestampFile(testFileDir); + + ValueType[] fullSchema = {ValueType.INT32, ValueType.INT64, ValueType.STRING}; + String[] fullNames = {"id", "ts", "name"}; + try { + new FrameReaderParquet().readFrameFromHDFS(file, fullSchema, fullNames, 3, 3); + Assert.fail("Expected rejection of the INT96 column"); + } + catch(IOException e) { + Assert.assertTrue(e.getMessage(), e.getMessage().contains("INT96")); + } + + ValueType[] subSchema = {ValueType.INT32, ValueType.STRING}; + String[] subNames = {"id", "name"}; + FrameBlock expected = ParquetTestUtils.sparkReadAsFrame(file, subSchema, subNames); + FrameBlock actual = new FrameReaderParquet().readFrameFromHDFS(file, subSchema, subNames, 3, 2); + TestUtils.compareFrames(expected, actual, false); + } + + @Test + public void testEmptyFile() throws Exception { + File temp = Files.createTempFile("systemds_empty_parquet", ".parquet").toFile(); + try { + ValueType[] schema = {ValueType.INT64, ValueType.STRING}; + String[] names = {"a", "b"}; + FrameBlock empty = new FrameBlock(schema, names, new String[0][]); + new FrameWriterParquet().writeFrameToHDFS(empty, temp.getPath(), 0, 2); + + FrameBlock result = new FrameReaderParquet().readFrameFromHDFS(temp.getPath(), schema, names, 0, 2); + + Assert.assertEquals("Empty file should yield 0 rows", 0, result.getNumRows()); + Assert.assertEquals(2, result.getNumColumns()); + } + finally { + temp.delete(); + } + } + + @Test + public void testReadSparkOutputDirectory() throws Exception { + // raw Spark output layout: a directory with _SUCCESS marker, .crc files and uuid part names; + // repartition shuffles rows across the part files, so values are verified keyed by id + File dir = new File(testFileDir, "spark_dir_out"); + ValueType[] schema = {ValueType.INT64, ValueType.FP64, ValueType.STRING}; + String[] names = {"id", "val", "name"}; + SparkSession spark = ParquetTestUtils.sparkSession(); + spark.range(100).selectExpr("id", "id * 0.5d as val", "concat('r', id) as name").repartition(2).write() + .mode(SaveMode.Overwrite).parquet(dir.getPath()); + Assert.assertTrue(new File(dir, "_SUCCESS").exists()); + + FrameBlock serial = new FrameReaderParquet().readFrameFromHDFS(dir.getPath(), schema, names, 100, 3); + Assert.assertEquals(100, serial.getNumRows()); + boolean[] seen = new boolean[100]; + for(int r = 0; r < 100; r++) { + int id = ((Long) serial.get(r, 0)).intValue(); + Assert.assertFalse("duplicate id " + id, seen[id]); + seen[id] = true; + Assert.assertEquals(id * 0.5, (Double) serial.get(r, 1), 0); + Assert.assertEquals("r" + id, serial.get(r, 2)); + } + + FrameBlock parallel = new FrameReaderParquetParallel().readFrameFromHDFS(dir.getPath(), schema, names, 100, 3); + TestUtils.compareFrames(serial, parallel, false); + } + + @Test + public void testParallelReaderMatchesSequential() throws Exception { + for(String filename : FILENAMES) { + ParquetMetadataInfo info = ParquetTestUtils.inferMetadata(filename); + + FrameReaderParquet sequential = new FrameReaderParquet(); + FrameBlock expected = sequential.readFrameFromHDFS(filename, info.schema, info.names, info.rlen, info.clen); + + FrameReaderParquetParallel parallel = new FrameReaderParquetParallel(); + FrameBlock actual = parallel.readFrameFromHDFS(filename, info.schema, info.names, info.rlen, info.clen); + + TestUtils.compareFrames(expected, actual, false); + } + } +} diff --git a/src/test/java/org/apache/sysds/test/functions/io/parquet/WriteParquetTest.java b/src/test/java/org/apache/sysds/test/functions/io/parquet/WriteParquetTest.java new file mode 100644 index 00000000000..46434da26a2 --- /dev/null +++ b/src/test/java/org/apache/sysds/test/functions/io/parquet/WriteParquetTest.java @@ -0,0 +1,180 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.test.functions.io.parquet; + +import java.io.File; +import java.nio.file.Files; +import java.util.Arrays; +import java.util.Map; + +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.sysds.common.Types.ValueType; +import org.apache.sysds.runtime.frame.data.FrameBlock; +import org.apache.sysds.runtime.io.FrameReaderParquet; +import org.apache.sysds.runtime.io.FrameReaderParquetParallel; +import org.apache.sysds.runtime.io.FrameWriterParquet; +import org.apache.sysds.runtime.io.FrameWriterParquetParallel; +import org.apache.sysds.test.TestUtils; +import org.apache.sysds.test.functions.io.parquet.ParquetTestUtils.ParquetMetadataInfo; +import org.junit.After; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +/** + * Verifies the parquet frame writers against Spark's parquet reader. + * + * Column 0 carries the row index so Spark reads of multi-file outputs can be compared order-independently (sorted by + * id), since engines do not guarantee row order across files. + */ +public class WriteParquetTest { + + private static final String TEMP_FILE = System.getProperty("java.io.tmpdir") + + "/systemds_write_parquet_test.parquet"; + private static final String TEMP_PAR_PATH = System.getProperty("java.io.tmpdir") + + "/systemds_write_parquet_test_par"; + + private static final ValueType[] SCHEMA = {ValueType.INT64, ValueType.STRING, ValueType.FP64, ValueType.BOOLEAN, + ValueType.INT32, ValueType.FP32}; + private static final long SEED = 4669201; + + // See ParquetTestUtils.generatePublicTestFiles(): these are generated with Spark's DataFrameWriter + private static File testFileDir; + private static String[] PUBLIC_FILES; + + @BeforeClass + public static void generateTestFiles() throws Exception { + testFileDir = Files.createTempDirectory("systemds_parquet_public_test_files").toFile(); + Map files = ParquetTestUtils.generatePublicTestFiles(testFileDir); + PUBLIC_FILES = new String[] {files.get("userdata1"), files.get("alltypes_plain"), files.get("all")}; + } + + @AfterClass + public static void cleanupTestFiles() { + File[] children = testFileDir.listFiles(); + if(children != null) + for(File f : children) + f.delete(); + testFileDir.delete(); + } + + @After + public void cleanup() { + new File(TEMP_FILE).delete(); + deleteRecursive(new File(TEMP_PAR_PATH)); + } + + private static void deleteRecursive(File f) { + if(f.isDirectory()) + for(File c : f.listFiles()) + deleteRecursive(c); + f.delete(); + } + + @Test + public void testSparkReadsMultiPartFiles() throws Exception { + FrameBlock original = indexedRandomFrame(60); + writeTwoParts(original); + + FrameBlock result = sparkReadSorted(TEMP_PAR_PATH, original.getColumnNames()); + TestUtils.compareFrames(original, result, false); + } + + @Test + public void testMultiPartFileRoundtrip() throws Exception { + FrameBlock original = indexedRandomFrame(60); + writeTwoParts(original); + + FrameBlock result = new FrameReaderParquetParallel().readFrameFromHDFS(TEMP_PAR_PATH, SCHEMA, + original.getColumnNames(), 60, SCHEMA.length); + + TestUtils.compareFrames(original, result, false); + } + + @Test + public void testSequentialReaderMultiPartFileRoundtrip() throws Exception { + // the serial reader must also expand a directory into its part files, not just the parallel one + FrameBlock original = indexedRandomFrame(60); + writeTwoParts(original); + + FrameBlock result = new FrameReaderParquet().readFrameFromHDFS(TEMP_PAR_PATH, SCHEMA, original.getColumnNames(), + 60, SCHEMA.length); + + TestUtils.compareFrames(original, result, false); + } + + @Test + public void testForcedParallelWriterRoundTrip() throws Exception { + FrameBlock original = indexedRandomFrame(20); + + FrameWriterParquetParallel writer = new FrameWriterParquetParallel(); + writer.setForcedParallel(true); + writer.writeFrameToHDFS(original, TEMP_PAR_PATH, 20, SCHEMA.length); + + FrameBlock result = new FrameReaderParquetParallel().readFrameFromHDFS(TEMP_PAR_PATH, SCHEMA, + original.getColumnNames(), 20, SCHEMA.length); + + TestUtils.compareFrames(original, result, false); + } + + @Test + public void testRoundtripPublicFiles() throws Exception { + for(String filename : PUBLIC_FILES) { + ParquetMetadataInfo info = ParquetTestUtils.inferMetadata(filename); + + FrameReaderParquet reader = new FrameReaderParquet(); + FrameBlock original = reader.readFrameFromHDFS(filename, info.schema, info.names, info.rlen, info.clen); + + FrameWriterParquet writer = new FrameWriterParquet(); + writer.writeFrameToHDFS(original, TEMP_FILE, original.getNumRows(), original.getNumColumns()); + + FrameBlock result = reader.readFrameFromHDFS(TEMP_FILE, info.schema, info.names, info.rlen, info.clen); + + TestUtils.compareFrames(original, result, false); + + FrameBlock sparkResult = ParquetTestUtils.sparkReadAsFrame(TEMP_FILE, info.schema, info.names); + TestUtils.compareFrames(original, sparkResult, false); + } + } + + /** Random frame over all parquet-supported value types, with the row index in column 0. */ + private static FrameBlock indexedRandomFrame(int rows) { + FrameBlock fb = TestUtils.generateRandomFrameBlock(rows, SCHEMA, SEED); + for(int r = 0; r < rows; r++) + fb.set(r, 0, (long) r); + return fb; + } + + private static void writeTwoParts(FrameBlock fb) throws Exception { + int rows = fb.getNumRows(); + new File(TEMP_PAR_PATH).mkdir(); + FrameWriterParquet writer = new FrameWriterParquet(); + writer.writeFrameToHDFS(fb.slice(0, rows / 2 - 1), TEMP_PAR_PATH + "/part-0.parquet", rows / 2, SCHEMA.length); + writer.writeFrameToHDFS(fb.slice(rows / 2, rows - 1), TEMP_PAR_PATH + "/part-1.parquet", rows - rows / 2, + SCHEMA.length); + } + + private static FrameBlock sparkReadSorted(String path, String[] names) { + Dataset df = ParquetTestUtils.sparkSession().read().parquet(path) + .select(names[0], Arrays.copyOfRange(names, 1, names.length)).sort(names[0]); + return ParquetTestUtils.toFrameBlock(df.collectAsList(), SCHEMA, names); + } +}