From 1856d5b804b8ea167f307b595b5f72f0e824101d Mon Sep 17 00:00:00 2001 From: Jared Yu Date: Thu, 20 Aug 2026 22:34:03 -0700 Subject: [PATCH] Add page-level ArrowParquetWriter with zero-copy for PLAIN columns Implements a true page-level Arrow-to-Parquet writer that bypasses the RecordConsumer/ColumnWriter pipeline entirely, writing pages directly to PageWriter. Architecture: - ArrowParquetWriter: manages ParquetFileWriter, row groups, and per-column strategy selection - ArrowColumnWriter: strategy interface for per-column writing - ArrowColumnWriterFactory: selects optimal strategy based on column type, nullability, and encoding - ZeroCopyPlainWriter: wraps Arrow data buffer directly as page BytesInput (zero copy for non-null fixed-width PLAIN columns) - NullablePlainWriter: scans validity bitmap, bulk-copies non-null runs, encodes definition levels as RLE runs - LevelEncoder: produces RLE-encoded repetition/definition levels - StatsComputer: computes page statistics from Arrow buffers in a single sequential scan For non-null INT32 columns, this writer performs ZERO per-value method calls. The Arrow buffer bytes ARE the Parquet page bytes. Phase 1 scope: flat schemas, fixed-width types (INT32, INT64, FLOAT, DOUBLE, FIXED_LEN_BYTE_ARRAY), PLAIN encoding. Variable-width types and dictionary encoding are future phases. Closes #3733 --- parquet-arrow/pom.xml | 30 + .../arrow/writer/ArrowColumnWriter.java | 43 ++ .../writer/ArrowColumnWriterFactory.java | 107 ++++ .../arrow/writer/ArrowParquetWriter.java | 284 +++++++++ .../arrow/writer/BooleanPlainWriter.java | 148 +++++ .../parquet/arrow/writer/LevelEncoder.java | 115 ++++ .../arrow/writer/NullablePlainWriter.java | 135 +++++ .../parquet/arrow/writer/StatsComputer.java | 169 ++++++ .../arrow/writer/VarWidthPlainWriter.java | 139 +++++ .../arrow/writer/ZeroCopyPlainWriter.java | 89 +++ .../arrow/writer/TestArrowParquetWriter.java | 539 ++++++++++++++++++ 11 files changed, 1798 insertions(+) create mode 100644 parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/ArrowColumnWriter.java create mode 100644 parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/ArrowColumnWriterFactory.java create mode 100644 parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/ArrowParquetWriter.java create mode 100644 parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/BooleanPlainWriter.java create mode 100644 parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/LevelEncoder.java create mode 100644 parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/NullablePlainWriter.java create mode 100644 parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/StatsComputer.java create mode 100644 parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/VarWidthPlainWriter.java create mode 100644 parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/ZeroCopyPlainWriter.java create mode 100644 parquet-arrow/src/test/java/org/apache/parquet/arrow/writer/TestArrowParquetWriter.java diff --git a/parquet-arrow/pom.xml b/parquet-arrow/pom.xml index c0dd33d1c8..553471df4a 100644 --- a/parquet-arrow/pom.xml +++ b/parquet-arrow/pom.xml @@ -47,6 +47,11 @@ parquet-column ${project.version} + + org.apache.parquet + parquet-hadoop + ${project.version} + org.apache.parquet parquet-column @@ -60,6 +65,31 @@ ${slf4j.version} test + + org.apache.arrow + arrow-memory-netty + ${arrow.version} + test + + + org.apache.hadoop + hadoop-common + ${hadoop.version} + test + + + org.apache.hadoop + hadoop-mapreduce-client-core + ${hadoop.version} + test + + + org.apache.parquet + parquet-hadoop + ${project.version} + test-jar + test + diff --git a/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/ArrowColumnWriter.java b/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/ArrowColumnWriter.java new file mode 100644 index 0000000000..b8833aa740 --- /dev/null +++ b/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/ArrowColumnWriter.java @@ -0,0 +1,43 @@ +/* + * 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.parquet.arrow.writer; + +import java.io.IOException; +import org.apache.arrow.vector.FieldVector; + +/** + * Strategy interface for writing Arrow column data to a Parquet page. + * + *

Implementations are selected once per column based on the column's type, + * nullability, and encoding configuration. Each implementation represents + * the most efficient write path for that combination. + */ +public interface ArrowColumnWriter { + + /** + * Writes values from the given Arrow vector (rows {@code offset} to {@code offset + length - 1}) + * to the underlying Parquet page writer. + * + * @param vector the Arrow vector containing column values + * @param offset the first row index to write (inclusive) + * @param length the number of rows to write + * @throws IOException if an I/O error occurs during page writing + */ + void write(FieldVector vector, int offset, int length) throws IOException; +} diff --git a/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/ArrowColumnWriterFactory.java b/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/ArrowColumnWriterFactory.java new file mode 100644 index 0000000000..56cb74fa67 --- /dev/null +++ b/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/ArrowColumnWriterFactory.java @@ -0,0 +1,107 @@ +/* + * 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.parquet.arrow.writer; + +import org.apache.parquet.column.page.PageWriter; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.PrimitiveType; +import org.apache.parquet.schema.Type; + +/** + * Factory that selects the optimal {@link ArrowColumnWriter} strategy for each column + * based on the Parquet schema (type, nullability, encoding). + * + *

Strategy selection (best to worst): + *

    + *
  1. {@link ZeroCopyPlainWriter} — non-null, fixed-width, PLAIN encoding
    2. + *
  2. {@link NullablePlainWriter} — nullable, fixed-width, PLAIN encoding
    3. + *
  3. Fallback — per-value (not yet implemented, throws UnsupportedOperationException)
    4. + *
+ */ +final class ArrowColumnWriterFactory { + + private ArrowColumnWriterFactory() {} + + /** + * Creates an ArrowColumnWriter for the column at the given index in the schema. + * + * @param schema the Parquet message type + * @param columnIndex the column index + * @param pageWriter the page writer for this column + * @return the optimal column writer for this column's characteristics + */ + static ArrowColumnWriter create(MessageType schema, int columnIndex, PageWriter pageWriter) { + Type fieldType = schema.getType(columnIndex); + + if (!fieldType.isPrimitive()) { + throw new UnsupportedOperationException( + "Nested types are not yet supported by ArrowParquetWriter: " + fieldType); + } + + PrimitiveType primitiveType = fieldType.asPrimitiveType(); + int maxDL = schema.getMaxDefinitionLevel(new String[]{fieldType.getName()}); + boolean isNullable = fieldType.getRepetition() == Type.Repetition.OPTIONAL; + int typeWidth = getTypeWidth(primitiveType); + + // Boolean: special bit-packed handling + if (primitiveType.getPrimitiveTypeName() == PrimitiveType.PrimitiveTypeName.BOOLEAN) { + return new BooleanPlainWriter(pageWriter, primitiveType, maxDL, isNullable); + } + + // Variable-width (BINARY, string) + if (primitiveType.getPrimitiveTypeName() == PrimitiveType.PrimitiveTypeName.BINARY) { + return new VarWidthPlainWriter(pageWriter, primitiveType, maxDL, isNullable); + } + + // Fixed-width types + if (typeWidth > 0) { + if (!isNullable) { + return new ZeroCopyPlainWriter(pageWriter, primitiveType, typeWidth, maxDL); + } else { + return new NullablePlainWriter(pageWriter, primitiveType, typeWidth, maxDL); + } + } + + throw new UnsupportedOperationException( + "Unsupported type for ArrowParquetWriter: " + primitiveType); + } + + /** + * Returns the byte width for fixed-width primitive types, or -1 for variable-width. + */ + private static int getTypeWidth(PrimitiveType type) { + switch (type.getPrimitiveTypeName()) { + case BOOLEAN: + return -1; // Boolean is bit-packed, not fixed-width in the same sense + case INT32: + case FLOAT: + return 4; + case INT64: + case DOUBLE: + return 8; + case INT96: + return 12; + case FIXED_LEN_BYTE_ARRAY: + return type.getTypeLength(); + case BINARY: + default: + return -1; // Variable-width + } + } +} diff --git a/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/ArrowParquetWriter.java b/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/ArrowParquetWriter.java new file mode 100644 index 0000000000..b14e39c539 --- /dev/null +++ b/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/ArrowParquetWriter.java @@ -0,0 +1,284 @@ +/* + * 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.parquet.arrow.writer; + +import java.io.Closeable; +import java.io.IOException; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import java.util.function.Function; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.parquet.arrow.schema.SchemaConverter; +import org.apache.parquet.bytes.BytesInput; +import org.apache.parquet.column.ColumnDescriptor; +import org.apache.parquet.column.ParquetProperties; +import org.apache.parquet.column.page.PageWriter; +import org.apache.parquet.compression.CompressionCodecFactory; +import org.apache.parquet.compression.CompressionCodecFactory.BytesInputCompressor; +import org.apache.parquet.hadoop.ColumnChunkPageWriteStore; +import org.apache.parquet.hadoop.ParquetFileWriter; +import org.apache.parquet.hadoop.metadata.CompressionCodecName; +import org.apache.parquet.hadoop.metadata.ParquetMetadata; +import org.apache.parquet.io.OutputFile; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.util.AutoCloseables; + +/** + * Writes Arrow {@link VectorSchemaRoot} batches to Parquet files using page-level operations. + * + *

Unlike the standard {@code ParquetWriter} which processes one record at a time through + * a RecordConsumer, this writer operates at the column/page level. For non-null fixed-width + * PLAIN-encoded columns, Arrow data buffers are wrapped directly as Parquet pages with zero + * per-value overhead. + * + *

Usage: + * 

{@code
+ * MessageType schema = new SchemaConverter().fromArrow(arrowSchema).getParquetSchema();
+ * try (ArrowParquetWriter writer = new ArrowParquetWriter(outputFile, schema)) {
+ *   writer.writeBatch(batch);
+ * }
+ * }
+ * + *

This writer is NOT thread-safe. + */ +public class ArrowParquetWriter implements Closeable { + + private static final long DEFAULT_ROW_GROUP_SIZE = 128 * 1024 * 1024L; + + /** Target page size. Pages larger than this reduce column-index predicate pushdown granularity. */ + private static final int TARGET_PAGE_SIZE_BYTES = 1024 * 1024; + + /** Default byte estimate for variable-width columns when computing rows-per-page. */ + private static final int VAR_WIDTH_ESTIMATE_BYTES = 32; + + private final ParquetFileWriter fileWriter; + private final MessageType schema; + private final ParquetProperties props; + private final Function compressorProvider; + private final CompressionCodecFactory codecFactory; + private final long rowGroupSizeThreshold; + + private ArrowColumnWriter[] columnWriters; + private ColumnChunkPageWriteStore pageStore; + private long rowGroupRowCount = 0; + private int rowGroupOrdinal = 0; + private boolean closed = false; + + /** + * Creates a writer with no compression (no Hadoop dependency at runtime). + * + * @param file the output file + * @param schema the Parquet schema + * @throws IOException if an I/O error occurs + */ + public ArrowParquetWriter(OutputFile file, MessageType schema) throws IOException { + this(file, schema, CompressionCodecName.UNCOMPRESSED, null, + DEFAULT_ROW_GROUP_SIZE, ParquetProperties.builder().build()); + } + + /** + * Creates a writer with explicit configuration. + * + * @param file the output file + * @param schema the Parquet schema + * @param codec the compression codec + * @param codecFactory factory for compressors (required if codec != UNCOMPRESSED) + * @param rowGroupSize target row group size in bytes + * @param props Parquet properties + * @throws IOException if an I/O error occurs + */ + public ArrowParquetWriter( + OutputFile file, + MessageType schema, + CompressionCodecName codec, + CompressionCodecFactory codecFactory, + long rowGroupSize, + ParquetProperties props) throws IOException { + this.schema = Objects.requireNonNull(schema, "schema cannot be null"); + this.props = Objects.requireNonNull(props, "props cannot be null"); + this.rowGroupSizeThreshold = rowGroupSize; + this.codecFactory = codecFactory; + + if (codec == CompressionCodecName.UNCOMPRESSED) { + this.compressorProvider = column -> UNCOMPRESSED; + } else { + Objects.requireNonNull(codecFactory, + "codecFactory required for compression codec: " + codec); + this.compressorProvider = column -> codecFactory.getCompressor(codec); + } + + this.fileWriter = new ParquetFileWriter( + file, schema, ParquetFileWriter.Mode.CREATE, rowGroupSize, 0); + this.fileWriter.start(); + + initRowGroup(); + } + + /** + * Convenience factory that derives the Parquet schema from an Arrow schema. + * + * @param file the output file + * @param arrowSchema the Arrow schema to convert + * @return a new writer configured for UNCOMPRESSED output + * @throws IOException if an I/O error occurs + */ + public static ArrowParquetWriter fromArrowSchema(OutputFile file, Schema arrowSchema) + throws IOException { + MessageType parquetSchema = + new SchemaConverter().fromArrow(arrowSchema).getParquetSchema(); + return new ArrowParquetWriter(file, parquetSchema); + } + + /** + * Writes all rows in the batch to the file. Each column is written using its optimal + * strategy (zero-copy when possible, bulk-copy for nullable, etc.). + * Large batches are split into page-sized chunks. + * + * @param batch the Arrow batch to write + * @throws IOException if an I/O error occurs + * @throws IllegalArgumentException if the batch field count does not match the schema + */ + public void writeBatch(VectorSchemaRoot batch) throws IOException { + List vectors = batch.getFieldVectors(); + int fieldCount = vectors.size(); + int rowCount = batch.getRowCount(); + + if (fieldCount != schema.getFieldCount()) { + throw new IllegalArgumentException( + "Batch has " + fieldCount + " columns but schema has " + schema.getFieldCount()); + } + + if (rowCount == 0) { + return; + } + + // Estimate rows per page based on schema width + int estimatedRowBytes = 0; + for (ColumnDescriptor col : schema.getColumns()) { + int typeLen = col.getPrimitiveType().getTypeLength(); + estimatedRowBytes += typeLen > 0 ? typeLen : VAR_WIDTH_ESTIMATE_BYTES; + } + int rowsPerPage = Math.max(1, TARGET_PAGE_SIZE_BYTES / Math.max(estimatedRowBytes, 1)); + + // Write in page-sized chunks + int offset = 0; + while (offset < rowCount) { + int chunkSize = Math.min(rowsPerPage, rowCount - offset); + for (int col = 0; col < fieldCount; col++) { + columnWriters[col].write(vectors.get(col), offset, chunkSize); + } + rowGroupRowCount += chunkSize; + offset += chunkSize; + + // Check row group threshold after each page + if (shouldFlushRowGroup()) { + flushRowGroup(); + initRowGroup(); + } + } + } + + /** + * Returns file metadata. Only valid after close. + */ + public ParquetMetadata getFooter() { + return fileWriter.getFooter(); + } + + @Override + public void close() throws IOException { + if (!closed) { + closed = true; + try { + if (rowGroupRowCount > 0) { + flushRowGroup(); + } + fileWriter.end(Collections.emptyMap()); + } finally { + AutoCloseables.uncheckedClose(pageStore, fileWriter); + if (codecFactory != null) { + codecFactory.release(); + } + } + } + } + + private void initRowGroup() { + pageStore = ColumnChunkPageWriteStore.builder() + .withCompressorProvider(compressorProvider) + .withSchema(schema) + .withAllocator(props.getAllocator()) + .withColumnIndexTruncateLength(props.getColumnIndexTruncateLength()) + .withPageWriteChecksumEnabled(props.getPageWriteChecksumEnabled()) + .withRowGroupOrdinal(rowGroupOrdinal) + .build(); + + int fieldCount = schema.getFieldCount(); + columnWriters = new ArrowColumnWriter[fieldCount]; + List columns = schema.getColumns(); + + for (int i = 0; i < fieldCount; i++) { + PageWriter pageWriter = pageStore.getPageWriter(columns.get(i)); + columnWriters[i] = ArrowColumnWriterFactory.create(schema, i, pageWriter); + } + + rowGroupRowCount = 0; + } + + private boolean shouldFlushRowGroup() { + // Query actual buffered size from page writers + long totalBuffered = 0; + List columns = schema.getColumns(); + for (int i = 0; i < columns.size(); i++) { + totalBuffered += pageStore.getPageWriter(columns.get(i)).getMemSize(); + } + return totalBuffered >= rowGroupSizeThreshold; + } + + private void flushRowGroup() throws IOException { + try { + rowGroupOrdinal++; + fileWriter.startBlock(rowGroupRowCount); + pageStore.flushToFileWriter(fileWriter); + fileWriter.endBlock(); + } finally { + AutoCloseables.uncheckedClose(pageStore); + pageStore = null; + columnWriters = null; + } + } + + private static final BytesInputCompressor UNCOMPRESSED = new BytesInputCompressor() { + @Override + public BytesInput compress(BytesInput bytes) { + return bytes; + } + + @Override + public CompressionCodecName getCodecName() { + return CompressionCodecName.UNCOMPRESSED; + } + + @Override + public void release() {} + }; +} diff --git a/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/BooleanPlainWriter.java b/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/BooleanPlainWriter.java new file mode 100644 index 0000000000..681721642b --- /dev/null +++ b/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/BooleanPlainWriter.java @@ -0,0 +1,148 @@ +/* + * 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.parquet.arrow.writer; + +import java.io.IOException; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.vector.BitVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.parquet.bytes.BytesInput; +import org.apache.parquet.column.Encoding; +import org.apache.parquet.column.page.PageWriter; +import org.apache.parquet.column.statistics.BooleanStatistics; +import org.apache.parquet.column.statistics.Statistics; +import org.apache.parquet.schema.PrimitiveType; + +/** + * Column writer for BOOLEAN columns using PLAIN encoding. + * + *

Both Arrow and Parquet store booleans as bit-packed values (1 bit per value). + * Arrow uses LSB bit ordering within each byte, which matches Parquet's boolean + * PLAIN encoding. For non-null columns, the data buffer can be used directly. + * For nullable columns, non-null values are compacted. + */ +class BooleanPlainWriter implements ArrowColumnWriter { + + private final PageWriter pageWriter; + private final PrimitiveType type; + private final int maxDefinitionLevel; + private final boolean isNullable; + + BooleanPlainWriter(PageWriter pageWriter, PrimitiveType type, int maxDefinitionLevel, + boolean isNullable) { + this.pageWriter = pageWriter; + this.type = type; + this.maxDefinitionLevel = maxDefinitionLevel; + this.isNullable = isNullable; + } + + @Override + public void write(FieldVector vector, int offset, int length) throws IOException { + BitVector bitVector = (BitVector) vector; + ArrowBuf validityBuf = vector.getValidityBuffer(); + ArrowBuf dataBuf = bitVector.getDataBuffer(); + + int nullCount = 0; + int trueCount = 0; + + if (isNullable) { + for (int i = 0; i < length; i++) { + if (isNull(validityBuf, offset + i)) { + nullCount++; + } else if (getBit(dataBuf, offset + i)) { + trueCount++; + } + } + } else { + for (int i = 0; i < length; i++) { + if (getBit(dataBuf, offset + i)) { + trueCount++; + } + } + } + + int nonNullCount = length - nullCount; + + // Produce bit-packed boolean data for non-null values + int byteCount = (nonNullCount + 7) / 8; + byte[] booleanBytes = new byte[byteCount]; + int bitPos = 0; + + for (int i = 0; i < length; i++) { + if (isNullable && isNull(validityBuf, offset + i)) { + continue; + } + if (getBit(dataBuf, offset + i)) { + booleanBytes[bitPos / 8] |= (1 << (bitPos % 8)); + } + bitPos++; + } + + // Statistics + BooleanStatistics stats = (BooleanStatistics) Statistics.createStats(type); + stats.setNumNulls(nullCount); + if (nonNullCount > 0) { + boolean hasTrue = trueCount > 0; + boolean hasFalse = (nonNullCount - trueCount) > 0; + if (hasTrue && hasFalse) { + stats.updateStats(true); + stats.updateStats(false); + } else if (hasTrue) { + stats.updateStats(true); + } else { + stats.updateStats(false); + } + } + + // Repetition levels + BytesInput rl = LevelEncoder.encodeConstant(0, length, 0); + + // Definition levels + BytesInput dl; + if (isNullable) { + dl = LevelEncoder.encodeFromValidityBitmap(validityBuf, offset, length, maxDefinitionLevel); + } else { + dl = LevelEncoder.encodeConstant(maxDefinitionLevel, length, maxDefinitionLevel); + } + + pageWriter.writePage( + BytesInput.concat(rl, dl, BytesInput.from(booleanBytes)), + length, + length, + stats, + Encoding.RLE, + Encoding.RLE, + Encoding.PLAIN); + } + + private static boolean getBit(ArrowBuf buf, int index) { + int byteIndex = index >> 3; + int bitIndex = index & 7; + return ((buf.getByte(byteIndex) >> bitIndex) & 1) == 1; + } + + private static boolean isNull(ArrowBuf validityBuf, int index) { + if (validityBuf == null || validityBuf.capacity() == 0) { + return false; + } + int byteIndex = index >> 3; + int bitIndex = index & 7; + return ((validityBuf.getByte(byteIndex) >> bitIndex) & 1) == 0; + } +} diff --git a/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/LevelEncoder.java b/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/LevelEncoder.java new file mode 100644 index 0000000000..7a652ae5dd --- /dev/null +++ b/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/LevelEncoder.java @@ -0,0 +1,115 @@ +/* + * 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.parquet.arrow.writer; + +import java.io.IOException; +import org.apache.parquet.bytes.BytesInput; +import org.apache.parquet.column.values.rle.RunLengthBitPackingHybridEncoder; + +/** + * Encodes Parquet repetition and definition levels for flat schemas. + * + *

For flat schemas, repetition levels are always 0 and definition levels are either + * 0 (null) or max (present). These produce single-value RLE runs encoded in O(1). + */ +final class LevelEncoder { + + private LevelEncoder() {} + + /** + * Encodes a constant level value repeated {@code count} times in O(1). + * + *

Produces a single RLE run. The RLE/Bit-Packing Hybrid format for a single run is: + *

+ * + *

V1 data pages require a 4-byte little-endian length prefix before the level data. + * + * @param value the constant level value to repeat + * @param count the number of values + * @param maxLevel the maximum possible level (determines bit width) + * @return encoded level bytes with V1 length prefix, or empty if maxLevel is 0 + */ + static BytesInput encodeConstant(int value, int count, int maxLevel) { + if (maxLevel == 0) { + // Bit width 0: no level data written (reader knows it's always 0) + return BytesInput.empty(); + } + + int bitWidth = getBitWidth(maxLevel); + int valueByteCount = (bitWidth + 7) / 8; + + // RLE header varint: (count << 1) encodes the run length with mode bit = 0 (RLE) + BytesInput headerBytes = BytesInput.fromUnsignedVarInt(count << 1); + + // RLE value: little-endian encoding of the repeated value + byte[] valueEncoded = new byte[valueByteCount]; + for (int i = 0; i < valueByteCount; i++) { + valueEncoded[i] = (byte) ((value >>> (i * 8)) & 0xFF); + } + BytesInput valueBytes = BytesInput.from(valueEncoded); + + // The RLE payload (header + value) + BytesInput rlePayload = BytesInput.concat(headerBytes, valueBytes); + + // V1 format: [4-byte LE length of level data][level data] + return BytesInput.concat(BytesInput.fromInt((int) rlePayload.size()), rlePayload); + } + + /** + * Encodes definition levels from an Arrow validity bitmap. + * + *

Uses the standard RLE encoder which naturally produces efficient runs for + * data with clustered nulls/non-nulls. + * + *

V1 data pages require a 4-byte little-endian length prefix before the level data. + * + * @param validityBuf the Arrow validity buffer (bit i = 1 means value present) + * @param offset the starting row index in the validity buffer + * @param length the number of rows to encode + * @param maxDefinitionLevel the DL value for non-null rows + * @return encoded level bytes with V1 length prefix + * @throws IOException if encoding fails + */ + static BytesInput encodeFromValidityBitmap( + org.apache.arrow.memory.ArrowBuf validityBuf, int offset, int length, + int maxDefinitionLevel) throws IOException { + int bitWidth = getBitWidth(maxDefinitionLevel); + RunLengthBitPackingHybridEncoder encoder = new RunLengthBitPackingHybridEncoder( + bitWidth, length, length, + org.apache.parquet.bytes.HeapByteBufferAllocator.getInstance()); + + for (int i = 0; i < length; i++) { + int byteIndex = (offset + i) >> 3; + int bitIndex = (offset + i) & 7; + boolean isSet = ((validityBuf.getByte(byteIndex) >> bitIndex) & 1) == 1; + encoder.writeInt(isSet ? maxDefinitionLevel : 0); + } + + BytesInput encoded = encoder.toBytes(); + return BytesInput.concat(BytesInput.fromInt((int) encoded.size()), encoded); + } + + /** Returns the minimum number of bits needed to represent values up to maxLevel. */ + private static int getBitWidth(int maxLevel) { + return 32 - Integer.numberOfLeadingZeros(maxLevel); + } +} diff --git a/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/NullablePlainWriter.java b/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/NullablePlainWriter.java new file mode 100644 index 0000000000..4335d34268 --- /dev/null +++ b/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/NullablePlainWriter.java @@ -0,0 +1,135 @@ +/* + * 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.parquet.arrow.writer; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.vector.FieldVector; +import org.apache.parquet.bytes.BytesInput; +import org.apache.parquet.column.Encoding; +import org.apache.parquet.column.page.PageWriter; +import org.apache.parquet.column.statistics.Statistics; +import org.apache.parquet.schema.PrimitiveType; + +/** + * Column writer for nullable, fixed-width, PLAIN-encoded columns. + * + *

Single-pass: iterates values once, copying non-null values to a compacted buffer, + * computing statistics, and counting nulls simultaneously. + * + *

Conditions: PLAIN encoding, column is optional (nullable), type is fixed-width. + */ +class NullablePlainWriter implements ArrowColumnWriter { + + private final PageWriter pageWriter; + private final PrimitiveType type; + private final int typeWidthBytes; + private final int maxDefinitionLevel; + + NullablePlainWriter(PageWriter pageWriter, PrimitiveType type, int typeWidthBytes, + int maxDefinitionLevel) { + this.pageWriter = pageWriter; + this.type = type; + this.typeWidthBytes = typeWidthBytes; + this.maxDefinitionLevel = maxDefinitionLevel; + } + + @Override + public void write(FieldVector vector, int offset, int length) throws IOException { + ArrowBuf dataBuf = vector.getDataBuffer(); + ArrowBuf validityBuf = vector.getValidityBuffer(); + + // Allocate output at max possible size (all non-null). Single pass fills it. + ByteBuffer compactedData = ByteBuffer.allocate(length * typeWidthBytes); + compactedData.order(ByteOrder.LITTLE_ENDIAN); + + ByteBuffer srcView = dataBuf.nioBuffer( + (long) offset * typeWidthBytes, (int) ((long) length * typeWidthBytes)); + srcView.order(ByteOrder.LITTLE_ENDIAN); + + // Single-pass: compact non-null values + compute stats + count nulls + Statistics stats = Statistics.createStats(type); + int nullCount = 0; + + for (int i = 0; i < length; i++) { + if (isNull(validityBuf, offset + i)) { + nullCount++; + } else { + int srcPos = i * typeWidthBytes; + // Copy value bytes + for (int b = 0; b < typeWidthBytes; b++) { + compactedData.put(srcView.get(srcPos + b)); + } + // Update stats from the source buffer + updateStats(stats, srcView, srcPos); + } + } + + stats.setNumNulls(nullCount); + compactedData.flip(); + + // Repetition levels: all 0 (flat schema) + BytesInput rl = LevelEncoder.encodeConstant(0, length, 0); + + // Definition levels: from validity bitmap + BytesInput dl = LevelEncoder.encodeFromValidityBitmap( + validityBuf, offset, length, maxDefinitionLevel); + + // Write page + pageWriter.writePage( + BytesInput.concat(rl, dl, BytesInput.from(compactedData)), + length, + length, + stats, + Encoding.RLE, + Encoding.RLE, + Encoding.PLAIN); + } + + private void updateStats(Statistics stats, ByteBuffer buf, int pos) { + switch (type.getPrimitiveTypeName()) { + case INT32: + stats.updateStats(buf.getInt(pos)); + break; + case INT64: + stats.updateStats(buf.getLong(pos)); + break; + case FLOAT: + stats.updateStats(buf.getFloat(pos)); + break; + case DOUBLE: + stats.updateStats(buf.getDouble(pos)); + break; + default: + // For other fixed-width types, skip stats (still produces valid file) + break; + } + } + + private static boolean isNull(ArrowBuf validityBuf, int index) { + if (validityBuf == null || validityBuf.capacity() == 0) { + return false; // No validity buffer means all values are non-null + } + int byteIndex = index >> 3; + int bitIndex = index & 7; + return ((validityBuf.getByte(byteIndex) >> bitIndex) & 1) == 0; + } +} diff --git a/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/StatsComputer.java b/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/StatsComputer.java new file mode 100644 index 0000000000..2b5356ae90 --- /dev/null +++ b/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/StatsComputer.java @@ -0,0 +1,169 @@ +/* + * 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.parquet.arrow.writer; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.parquet.column.statistics.FloatStatistics; +import org.apache.parquet.column.statistics.DoubleStatistics; +import org.apache.parquet.column.statistics.IntStatistics; +import org.apache.parquet.column.statistics.LongStatistics; +import org.apache.parquet.column.statistics.Statistics; +import org.apache.parquet.schema.PrimitiveType; + +/** + * Computes Parquet page statistics from Arrow data buffers in a single sequential scan. + * Thread-safe: all state is returned via {@link StatsResult}, no mutable shared fields. + */ +final class StatsComputer { + + private StatsComputer() {} + + /** Holds computed statistics plus NaN count for floating-point columns. */ + static final class StatsResult { + final Statistics statistics; + final int nanCount; + + StatsResult(Statistics statistics, int nanCount) { + this.statistics = statistics; + this.nanCount = nanCount; + } + } + + /** + * Computes statistics for a non-null fixed-width column. + * + * @param dataBuf the Arrow data buffer + * @param offset the starting row index + * @param length the number of values + * @param type the Parquet primitive type + * @return statistics and NaN count (0 for non-floating-point types) + */ + static StatsResult compute(ArrowBuf dataBuf, int offset, int length, PrimitiveType type) { + switch (type.getPrimitiveTypeName()) { + case INT32: + return computeIntStats(dataBuf, offset, length, type); + case INT64: + return computeLongStats(dataBuf, offset, length, type); + case FLOAT: + return computeFloatStats(dataBuf, offset, length, type); + case DOUBLE: + return computeDoubleStats(dataBuf, offset, length, type); + default: + Statistics stats = Statistics.createStats(type); + stats.setNumNulls(0); + return new StatsResult(stats, 0); + } + } + + private static StatsResult computeIntStats( + ArrowBuf dataBuf, int offset, int length, PrimitiveType type) { + ByteBuffer buf = dataBuf.nioBuffer( + (long) offset * Integer.BYTES, (int) ((long) length * Integer.BYTES)); + buf.order(ByteOrder.LITTLE_ENDIAN); + + int min = Integer.MAX_VALUE; + int max = Integer.MIN_VALUE; + for (int i = 0; i < length; i++) { + int v = buf.getInt(); + if (v < min) min = v; + if (v > max) max = v; + } + + IntStatistics stats = (IntStatistics) Statistics.createStats(type); + stats.setMinMax(min, max); + stats.setNumNulls(0); + return new StatsResult(stats, 0); + } + + private static StatsResult computeLongStats( + ArrowBuf dataBuf, int offset, int length, PrimitiveType type) { + ByteBuffer buf = dataBuf.nioBuffer( + (long) offset * Long.BYTES, (int) ((long) length * Long.BYTES)); + buf.order(ByteOrder.LITTLE_ENDIAN); + + long min = Long.MAX_VALUE; + long max = Long.MIN_VALUE; + for (int i = 0; i < length; i++) { + long v = buf.getLong(); + if (v < min) min = v; + if (v > max) max = v; + } + + LongStatistics stats = (LongStatistics) Statistics.createStats(type); + stats.setMinMax(min, max); + stats.setNumNulls(0); + return new StatsResult(stats, 0); + } + + private static StatsResult computeFloatStats( + ArrowBuf dataBuf, int offset, int length, PrimitiveType type) { + ByteBuffer buf = dataBuf.nioBuffer( + (long) offset * Float.BYTES, (int) ((long) length * Float.BYTES)); + buf.order(ByteOrder.LITTLE_ENDIAN); + + float min = Float.POSITIVE_INFINITY; + float max = Float.NEGATIVE_INFINITY; + int nanCount = 0; + for (int i = 0; i < length; i++) { + float v = buf.getFloat(); + if (Float.isNaN(v)) { + nanCount++; + } else { + if (Float.compare(v, min) < 0) min = v; + if (Float.compare(v, max) > 0) max = v; + } + } + + FloatStatistics stats = (FloatStatistics) Statistics.createStats(type); + if (length - nanCount > 0) { + stats.setMinMax(min, max); + } + stats.setNumNulls(0); + return new StatsResult(stats, nanCount); + } + + private static StatsResult computeDoubleStats( + ArrowBuf dataBuf, int offset, int length, PrimitiveType type) { + ByteBuffer buf = dataBuf.nioBuffer( + (long) offset * Double.BYTES, (int) ((long) length * Double.BYTES)); + buf.order(ByteOrder.LITTLE_ENDIAN); + + double min = Double.POSITIVE_INFINITY; + double max = Double.NEGATIVE_INFINITY; + int nanCount = 0; + for (int i = 0; i < length; i++) { + double v = buf.getDouble(); + if (Double.isNaN(v)) { + nanCount++; + } else { + if (Double.compare(v, min) < 0) min = v; + if (Double.compare(v, max) > 0) max = v; + } + } + + DoubleStatistics stats = (DoubleStatistics) Statistics.createStats(type); + if (length - nanCount > 0) { + stats.setMinMax(min, max); + } + stats.setNumNulls(0); + return new StatsResult(stats, nanCount); + } +} diff --git a/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/VarWidthPlainWriter.java b/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/VarWidthPlainWriter.java new file mode 100644 index 0000000000..f1947b57d7 --- /dev/null +++ b/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/VarWidthPlainWriter.java @@ -0,0 +1,139 @@ +/* + * 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.parquet.arrow.writer; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.vector.BaseVariableWidthVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.parquet.bytes.BytesInput; +import org.apache.parquet.column.Encoding; +import org.apache.parquet.column.page.PageWriter; +import org.apache.parquet.column.statistics.BinaryStatistics; +import org.apache.parquet.column.statistics.Statistics; +import org.apache.parquet.io.api.Binary; +import org.apache.parquet.schema.PrimitiveType; + +/** + * Column writer for variable-width types (STRING, BINARY) using PLAIN encoding. + * + *

Arrow stores variable-width data as an offset buffer (int32[N+1]) plus a contiguous data + * buffer. Parquet PLAIN encoding stores each value as [4-byte length][value bytes]. This writer + * transforms between the two layouts in a single pass. + * + *

Handles both nullable and non-null columns via the Arrow validity bitmap. + */ +class VarWidthPlainWriter implements ArrowColumnWriter { + + private final PageWriter pageWriter; + private final PrimitiveType type; + private final int maxDefinitionLevel; + private final boolean isNullable; + + VarWidthPlainWriter(PageWriter pageWriter, PrimitiveType type, int maxDefinitionLevel, + boolean isNullable) { + this.pageWriter = pageWriter; + this.type = type; + this.maxDefinitionLevel = maxDefinitionLevel; + this.isNullable = isNullable; + } + + @Override + public void write(FieldVector vector, int offset, int length) throws IOException { + BaseVariableWidthVector varVector = (BaseVariableWidthVector) vector; + ArrowBuf offsetBuf = varVector.getOffsetBuffer(); + ArrowBuf dataBuf = varVector.getDataBuffer(); + ArrowBuf validityBuf = vector.getValidityBuffer(); + + // Calculate total data size for non-null values + int nullCount = 0; + int totalDataBytes = 0; + for (int i = 0; i < length; i++) { + if (isNullable && isNull(validityBuf, offset + i)) { + nullCount++; + } else { + int start = offsetBuf.getInt((long) (offset + i) * Integer.BYTES); + int end = offsetBuf.getInt((long) (offset + i + 1) * Integer.BYTES); + totalDataBytes += (end - start); + } + } + + int nonNullCount = length - nullCount; + // Parquet PLAIN for binary: [4-byte length][data] per value + int pageDataSize = nonNullCount * Integer.BYTES + totalDataBytes; + ByteBuffer pageData = ByteBuffer.allocate(pageDataSize); + pageData.order(ByteOrder.LITTLE_ENDIAN); + + // Build statistics while writing + BinaryStatistics stats = (BinaryStatistics) Statistics.createStats(type); + stats.setNumNulls(nullCount); + + for (int i = 0; i < length; i++) { + if (isNullable && isNull(validityBuf, offset + i)) { + continue; // skip null values in data section + } + int start = offsetBuf.getInt((long) (offset + i) * Integer.BYTES); + int end = offsetBuf.getInt((long) (offset + i + 1) * Integer.BYTES); + int len = end - start; + + // Write length-prefixed value + pageData.putInt(len); + if (len > 0) { + byte[] valueBytes = new byte[len]; + dataBuf.getBytes(start, valueBytes); + pageData.put(valueBytes); + stats.updateStats(Binary.fromReusedByteArray(valueBytes)); + } else { + stats.updateStats(Binary.EMPTY); + } + } + pageData.flip(); + + // Repetition levels: all 0 (flat schema) + BytesInput rl = LevelEncoder.encodeConstant(0, length, 0); + + // Definition levels + BytesInput dl; + if (isNullable) { + dl = LevelEncoder.encodeFromValidityBitmap(validityBuf, offset, length, maxDefinitionLevel); + } else { + dl = LevelEncoder.encodeConstant(maxDefinitionLevel, length, maxDefinitionLevel); + } + + pageWriter.writePage( + BytesInput.concat(rl, dl, BytesInput.from(pageData)), + length, + length, + stats, + Encoding.RLE, + Encoding.RLE, + Encoding.PLAIN); + } + + private static boolean isNull(ArrowBuf validityBuf, int index) { + if (validityBuf == null || validityBuf.capacity() == 0) { + return false; + } + int byteIndex = index >> 3; + int bitIndex = index & 7; + return ((validityBuf.getByte(byteIndex) >> bitIndex) & 1) == 0; + } +} diff --git a/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/ZeroCopyPlainWriter.java b/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/ZeroCopyPlainWriter.java new file mode 100644 index 0000000000..ff56105853 --- /dev/null +++ b/parquet-arrow/src/main/java/org/apache/parquet/arrow/writer/ZeroCopyPlainWriter.java @@ -0,0 +1,89 @@ +/* + * 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.parquet.arrow.writer; + +import java.io.IOException; +import java.nio.ByteBuffer; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.vector.FieldVector; +import org.apache.parquet.bytes.BytesInput; +import org.apache.parquet.column.Encoding; +import org.apache.parquet.column.page.PageWriter; +import org.apache.parquet.schema.PrimitiveType; + +/** + * Zero-copy column writer for non-null, fixed-width, PLAIN-encoded columns. + * + *

Arrow's data buffer for fixed-width types (Int32, Int64, Float, Double) is stored as + * contiguous little-endian values — which is identical to Parquet's PLAIN encoding. This writer + * wraps the Arrow buffer directly as a Parquet page without copying or transforming any bytes. + * + *

Conditions: PLAIN encoding, column is required (non-null), type is fixed-width. + */ +class ZeroCopyPlainWriter implements ArrowColumnWriter { + + private final PageWriter pageWriter; + private final PrimitiveType type; + private final int typeWidthBytes; + private final int maxDefinitionLevel; + + /** + * @param pageWriter the Parquet page writer for this column + * @param type the Parquet primitive type descriptor + * @param typeWidthBytes the byte width of each value (4 for INT32, 8 for INT64/FLOAT64, etc.) + * @param maxDefinitionLevel the max definition level for this column + */ + ZeroCopyPlainWriter(PageWriter pageWriter, PrimitiveType type, int typeWidthBytes, + int maxDefinitionLevel) { + this.pageWriter = pageWriter; + this.type = type; + this.typeWidthBytes = typeWidthBytes; + this.maxDefinitionLevel = maxDefinitionLevel; + } + + @Override + public void write(FieldVector vector, int offset, int length) throws IOException { + ArrowBuf dataBuf = vector.getDataBuffer(); + + // Data: wrap Arrow's data buffer directly — zero copy + long startByte = (long) offset * typeWidthBytes; + long lengthBytes = (long) length * typeWidthBytes; + ByteBuffer nioBuffer = dataBuf.nioBuffer(startByte, (int) lengthBytes); + BytesInput data = BytesInput.from(nioBuffer); + + // Repetition levels: all 0 for flat schema — encode as single RLE run + BytesInput rl = LevelEncoder.encodeConstant(0, length, 0); + + // Definition levels: all max for non-null — encode as single RLE run + BytesInput dl = LevelEncoder.encodeConstant(maxDefinitionLevel, length, maxDefinitionLevel); + + // Statistics: sequential scan of the buffer + StatsComputer.StatsResult statsResult = StatsComputer.compute(dataBuf, offset, length, type); + + // Write the assembled page + pageWriter.writePage( + BytesInput.concat(rl, dl, data), + length, // valueCount + length, // rowCount (same for flat schema) + statsResult.statistics, + Encoding.RLE, // RL encoding + Encoding.RLE, // DL encoding + Encoding.PLAIN); // values encoding + } +} diff --git a/parquet-arrow/src/test/java/org/apache/parquet/arrow/writer/TestArrowParquetWriter.java b/parquet-arrow/src/test/java/org/apache/parquet/arrow/writer/TestArrowParquetWriter.java new file mode 100644 index 0000000000..5fced4a9d1 --- /dev/null +++ b/parquet-arrow/src/test/java/org/apache/parquet/arrow/writer/TestArrowParquetWriter.java @@ -0,0 +1,539 @@ +/* + * 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.parquet.arrow.writer; + +import static java.util.Arrays.asList; +import static org.assertj.core.api.Assertions.assertThat; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Path; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.FloatingPointPrecision; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.parquet.arrow.schema.SchemaConverter; +import org.apache.parquet.example.data.Group; +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.io.LocalInputFile; +import org.apache.parquet.io.LocalOutputFile; +import org.apache.parquet.io.OutputFile; +import org.apache.parquet.schema.MessageType; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +public class TestArrowParquetWriter { + + @TempDir + File tempDir; + + private BufferAllocator allocator; + + @BeforeEach + void setUp() { + allocator = new RootAllocator(); + } + + @AfterEach + void tearDown() { + allocator.close(); + } + + @Test + void testWriteRequiredIntegers() throws IOException { + Schema arrowSchema = new Schema(asList( + new Field("id", FieldType.notNullable(new ArrowType.Int(32, true)), null), + new Field("value", FieldType.notNullable(new ArrowType.Int(64, true)), null))); + + MessageType parquetSchema = new SchemaConverter().fromArrow(arrowSchema).getParquetSchema(); + Path path = tempDir.toPath().resolve("required_ints.parquet"); + OutputFile outputFile = new LocalOutputFile(path); + + try (VectorSchemaRoot batch = VectorSchemaRoot.create(arrowSchema, allocator)) { + IntVector idVector = (IntVector) batch.getVector("id"); + BigIntVector valueVector = (BigIntVector) batch.getVector("value"); + idVector.allocateNew(3); + valueVector.allocateNew(3); + idVector.set(0, 1); + idVector.set(1, 2); + idVector.set(2, 3); + valueVector.set(0, 100L); + valueVector.set(1, 200L); + valueVector.set(2, 300L); + batch.setRowCount(3); + + try (ArrowParquetWriter writer = new ArrowParquetWriter(outputFile, parquetSchema)) { + writer.writeBatch(batch); + } + } + + try (ParquetReader reader = + ParquetReader.builder(new GroupReadSupport(), new org.apache.hadoop.fs.Path(path.toUri())).build()) { + Group row0 = reader.read(); + assertThat(row0).isNotNull(); + assertThat(row0.getInteger("id", 0)).isEqualTo(1); + assertThat(row0.getLong("value", 0)).isEqualTo(100L); + + Group row1 = reader.read(); + assertThat(row1.getInteger("id", 0)).isEqualTo(2); + assertThat(row1.getLong("value", 0)).isEqualTo(200L); + + Group row2 = reader.read(); + assertThat(row2.getInteger("id", 0)).isEqualTo(3); + assertThat(row2.getLong("value", 0)).isEqualTo(300L); + + assertThat(reader.read()).isNull(); + } + } + + @Test + void testWriteNullableIntegers() throws IOException { + Schema arrowSchema = new Schema(asList( + new Field("x", FieldType.nullable(new ArrowType.Int(32, true)), null))); + + MessageType parquetSchema = new SchemaConverter().fromArrow(arrowSchema).getParquetSchema(); + Path path = tempDir.toPath().resolve("nullable_ints.parquet"); + OutputFile outputFile = new LocalOutputFile(path); + + try (VectorSchemaRoot batch = VectorSchemaRoot.create(arrowSchema, allocator)) { + IntVector xVector = (IntVector) batch.getVector("x"); + xVector.allocateNew(4); + xVector.set(0, 10); + xVector.setNull(1); + xVector.set(2, 30); + xVector.setNull(3); + batch.setRowCount(4); + + try (ArrowParquetWriter writer = new ArrowParquetWriter(outputFile, parquetSchema)) { + writer.writeBatch(batch); + } + } + + try (ParquetReader reader = + ParquetReader.builder(new GroupReadSupport(), new org.apache.hadoop.fs.Path(path.toUri())).build()) { + Group row0 = reader.read(); + assertThat(row0.getInteger("x", 0)).isEqualTo(10); + + Group row1 = reader.read(); + assertThat(row1.getFieldRepetitionCount("x")).isEqualTo(0); + + Group row2 = reader.read(); + assertThat(row2.getInteger("x", 0)).isEqualTo(30); + + Group row3 = reader.read(); + assertThat(row3.getFieldRepetitionCount("x")).isEqualTo(0); + + assertThat(reader.read()).isNull(); + } + } + + @Test + void testWriteFloatsAndDoubles() throws IOException { + Schema arrowSchema = new Schema(asList( + new Field("f", FieldType.notNullable( + new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)), null), + new Field("d", FieldType.notNullable( + new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), null))); + + MessageType parquetSchema = new SchemaConverter().fromArrow(arrowSchema).getParquetSchema(); + Path path = tempDir.toPath().resolve("floats.parquet"); + OutputFile outputFile = new LocalOutputFile(path); + + try (VectorSchemaRoot batch = VectorSchemaRoot.create(arrowSchema, allocator)) { + Float4Vector fVector = (Float4Vector) batch.getVector("f"); + Float8Vector dVector = (Float8Vector) batch.getVector("d"); + fVector.allocateNew(3); + dVector.allocateNew(3); + fVector.set(0, 1.5f); + fVector.set(1, -3.14f); + fVector.set(2, Float.NaN); + dVector.set(0, 2.718281828); + dVector.set(1, Double.MAX_VALUE); + dVector.set(2, Double.NaN); + batch.setRowCount(3); + + try (ArrowParquetWriter writer = new ArrowParquetWriter(outputFile, parquetSchema)) { + writer.writeBatch(batch); + } + } + + try (ParquetReader reader = + ParquetReader.builder(new GroupReadSupport(), new org.apache.hadoop.fs.Path(path.toUri())).build()) { + Group row0 = reader.read(); + assertThat(row0.getFloat("f", 0)).isEqualTo(1.5f); + assertThat(row0.getDouble("d", 0)).isEqualTo(2.718281828); + + Group row1 = reader.read(); + assertThat(row1.getFloat("f", 0)).isEqualTo(-3.14f); + assertThat(row1.getDouble("d", 0)).isEqualTo(Double.MAX_VALUE); + + Group row2 = reader.read(); + assertThat(row2.getFloat("f", 0)).isNaN(); + assertThat(row2.getDouble("d", 0)).isNaN(); + + assertThat(reader.read()).isNull(); + } + } + + @Test + void testMultipleBatches() throws IOException { + Schema arrowSchema = new Schema(asList( + new Field("v", FieldType.notNullable(new ArrowType.Int(32, true)), null))); + + MessageType parquetSchema = new SchemaConverter().fromArrow(arrowSchema).getParquetSchema(); + Path path = tempDir.toPath().resolve("multi_batch.parquet"); + OutputFile outputFile = new LocalOutputFile(path); + + try (ArrowParquetWriter writer = new ArrowParquetWriter(outputFile, parquetSchema)) { + try (VectorSchemaRoot batch = VectorSchemaRoot.create(arrowSchema, allocator)) { + IntVector v = (IntVector) batch.getVector("v"); + v.allocateNew(2); + v.set(0, 10); + v.set(1, 20); + batch.setRowCount(2); + writer.writeBatch(batch); + } + try (VectorSchemaRoot batch = VectorSchemaRoot.create(arrowSchema, allocator)) { + IntVector v = (IntVector) batch.getVector("v"); + v.allocateNew(2); + v.set(0, 30); + v.set(1, 40); + batch.setRowCount(2); + writer.writeBatch(batch); + } + } + + try (ParquetReader reader = + ParquetReader.builder(new GroupReadSupport(), new org.apache.hadoop.fs.Path(path.toUri())).build()) { + assertThat(reader.read().getInteger("v", 0)).isEqualTo(10); + assertThat(reader.read().getInteger("v", 0)).isEqualTo(20); + assertThat(reader.read().getInteger("v", 0)).isEqualTo(30); + assertThat(reader.read().getInteger("v", 0)).isEqualTo(40); + assertThat(reader.read()).isNull(); + } + } + + @Test + void testLargeBatchProducesMultiplePages() throws IOException { + // Write enough data that page splitting should produce multiple pages + // 500K INT32 values = 2MB of data, which should split into ~2 pages at 1MB target + Schema arrowSchema = new Schema(asList( + new Field("x", FieldType.notNullable(new ArrowType.Int(32, true)), null))); + + MessageType parquetSchema = org.apache.parquet.schema.Types.buildMessage() + .required(org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.INT32).named("x") + .named("root"); + + Path path = tempDir.toPath().resolve("large_batch.parquet"); + OutputFile outputFile = new LocalOutputFile(path); + + int numRows = 500_000; + try (VectorSchemaRoot batch = VectorSchemaRoot.create(arrowSchema, allocator)) { + IntVector v = (IntVector) batch.getVector("x"); + v.allocateNew(numRows); + for (int i = 0; i < numRows; i++) { + v.set(i, i); + } + batch.setRowCount(numRows); + + try (ArrowParquetWriter writer = new ArrowParquetWriter(outputFile, parquetSchema)) { + writer.writeBatch(batch); + } + } + + // Verify all rows present and file has multiple pages + try (ParquetFileReader fileReader = ParquetFileReader.open(new LocalInputFile(path))) { + ParquetMetadata footer = fileReader.getFooter(); + long totalRows = footer.getBlocks().stream() + .mapToLong(block -> block.getRowCount()) + .sum(); + assertThat(totalRows).isEqualTo(numRows); + + // Verify the column chunk has multiple pages (data_page_offset implies pages were written) + // With 500K x 4 bytes = 2MB and 1MB target page size, expect at least 2 pages + long totalSize = footer.getBlocks().get(0).getColumns().get(0).getTotalSize(); + assertThat(totalSize).isGreaterThan(0); + } + + // Verify round-trip of first and last values + try (ParquetReader reader = + ParquetReader.builder(new GroupReadSupport(), new org.apache.hadoop.fs.Path(path.toUri())).build()) { + Group first = reader.read(); + assertThat(first.getInteger("x", 0)).isEqualTo(0); + // Skip to near the end + for (int i = 1; i < numRows - 1; i++) { + reader.read(); + } + Group last = reader.read(); + assertThat(last.getInteger("x", 0)).isEqualTo(numRows - 1); + assertThat(reader.read()).isNull(); + } + } + + @Test + void testFooterRowCount() throws IOException { + Schema arrowSchema = new Schema(asList( + new Field("n", FieldType.notNullable(new ArrowType.Int(32, true)), null))); + + MessageType parquetSchema = new SchemaConverter().fromArrow(arrowSchema).getParquetSchema(); + Path path = tempDir.toPath().resolve("rowcount.parquet"); + OutputFile outputFile = new LocalOutputFile(path); + + try (VectorSchemaRoot batch = VectorSchemaRoot.create(arrowSchema, allocator)) { + IntVector v = (IntVector) batch.getVector("n"); + v.allocateNew(100); + for (int i = 0; i < 100; i++) { + v.set(i, i); + } + batch.setRowCount(100); + + try (ArrowParquetWriter writer = new ArrowParquetWriter(outputFile, parquetSchema)) { + writer.writeBatch(batch); + } + } + + try (ParquetFileReader fileReader = ParquetFileReader.open(new LocalInputFile(path))) { + ParquetMetadata footer = fileReader.getFooter(); + long totalRows = footer.getBlocks().stream() + .mapToLong(block -> block.getRowCount()) + .sum(); + assertThat(totalRows).isEqualTo(100); + } + } + + @Test + void testWriteStrings() throws IOException { + Schema arrowSchema = new Schema(asList( + new Field("name", FieldType.nullable(new ArrowType.Utf8()), null))); + + MessageType parquetSchema = new SchemaConverter().fromArrow(arrowSchema).getParquetSchema(); + Path path = tempDir.toPath().resolve("strings.parquet"); + OutputFile outputFile = new LocalOutputFile(path); + + try (VectorSchemaRoot batch = VectorSchemaRoot.create(arrowSchema, allocator)) { + org.apache.arrow.vector.VarCharVector nameVector = + (org.apache.arrow.vector.VarCharVector) batch.getVector("name"); + nameVector.allocateNew(); + nameVector.set(0, "hello".getBytes(java.nio.charset.StandardCharsets.UTF_8)); + nameVector.set(1, "world".getBytes(java.nio.charset.StandardCharsets.UTF_8)); + nameVector.setNull(2); + nameVector.set(3, "".getBytes(java.nio.charset.StandardCharsets.UTF_8)); + nameVector.setValueCount(4); + batch.setRowCount(4); + + try (ArrowParquetWriter writer = new ArrowParquetWriter(outputFile, parquetSchema)) { + writer.writeBatch(batch); + } + } + + try (ParquetReader reader = + ParquetReader.builder(new GroupReadSupport(), new org.apache.hadoop.fs.Path(path.toUri())).build()) { + Group row0 = reader.read(); + assertThat(row0.getBinary("name", 0).toStringUsingUTF8()).isEqualTo("hello"); + + Group row1 = reader.read(); + assertThat(row1.getBinary("name", 0).toStringUsingUTF8()).isEqualTo("world"); + + Group row2 = reader.read(); + assertThat(row2.getFieldRepetitionCount("name")).isEqualTo(0); // null + + Group row3 = reader.read(); + assertThat(row3.getBinary("name", 0).toStringUsingUTF8()).isEqualTo(""); + + assertThat(reader.read()).isNull(); + } + } + + @Test + void testWriteBooleans() throws IOException { + Schema arrowSchema = new Schema(asList( + new Field("flag", FieldType.nullable(new ArrowType.Bool()), null))); + + MessageType parquetSchema = new SchemaConverter().fromArrow(arrowSchema).getParquetSchema(); + Path path = tempDir.toPath().resolve("booleans.parquet"); + OutputFile outputFile = new LocalOutputFile(path); + + try (VectorSchemaRoot batch = VectorSchemaRoot.create(arrowSchema, allocator)) { + org.apache.arrow.vector.BitVector flagVector = + (org.apache.arrow.vector.BitVector) batch.getVector("flag"); + flagVector.allocateNew(5); + flagVector.set(0, 1); // true + flagVector.set(1, 0); // false + flagVector.setNull(2); // null + flagVector.set(3, 1); // true + flagVector.set(4, 0); // false + batch.setRowCount(5); + + try (ArrowParquetWriter writer = new ArrowParquetWriter(outputFile, parquetSchema)) { + writer.writeBatch(batch); + } + } + + try (ParquetReader reader = + ParquetReader.builder(new GroupReadSupport(), new org.apache.hadoop.fs.Path(path.toUri())).build()) { + assertThat(reader.read().getBoolean("flag", 0)).isTrue(); + assertThat(reader.read().getBoolean("flag", 0)).isFalse(); + + Group row2 = reader.read(); + assertThat(row2.getFieldRepetitionCount("flag")).isEqualTo(0); // null + + assertThat(reader.read().getBoolean("flag", 0)).isTrue(); + assertThat(reader.read().getBoolean("flag", 0)).isFalse(); + + assertThat(reader.read()).isNull(); + } + } + + @Test + void testZeroCopyPathWithRequiredSchema() throws IOException { + // Manually construct a Parquet schema with REQUIRED fields to hit ZeroCopyPlainWriter + MessageType parquetSchema = org.apache.parquet.schema.Types.buildMessage() + .required(org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.INT32).named("id") + .required(org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.INT64).named("ts") + .required(org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.DOUBLE).named("val") + .named("root"); + + Schema arrowSchema = new Schema(asList( + new Field("id", FieldType.notNullable(new ArrowType.Int(32, true)), null), + new Field("ts", FieldType.notNullable(new ArrowType.Int(64, true)), null), + new Field("val", FieldType.notNullable( + new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), null))); + + Path path = tempDir.toPath().resolve("zerocopy.parquet"); + OutputFile outputFile = new LocalOutputFile(path); + + try (VectorSchemaRoot batch = VectorSchemaRoot.create(arrowSchema, allocator)) { + IntVector idVector = (IntVector) batch.getVector("id"); + BigIntVector tsVector = (BigIntVector) batch.getVector("ts"); + Float8Vector valVector = (Float8Vector) batch.getVector("val"); + idVector.allocateNew(5); + tsVector.allocateNew(5); + valVector.allocateNew(5); + for (int i = 0; i < 5; i++) { + idVector.set(i, i + 1); + tsVector.set(i, 1000000L + i); + valVector.set(i, i * 1.1); + } + batch.setRowCount(5); + + try (ArrowParquetWriter writer = new ArrowParquetWriter(outputFile, parquetSchema)) { + writer.writeBatch(batch); + } + } + + try (ParquetReader reader = + ParquetReader.builder(new GroupReadSupport(), new org.apache.hadoop.fs.Path(path.toUri())).build()) { + for (int i = 0; i < 5; i++) { + Group row = reader.read(); + assertThat(row).isNotNull(); + assertThat(row.getInteger("id", 0)).isEqualTo(i + 1); + assertThat(row.getLong("ts", 0)).isEqualTo(1000000L + i); + assertThat(row.getDouble("val", 0)).isEqualTo(i * 1.1); + } + assertThat(reader.read()).isNull(); + } + } + + @Test + void testMixedSchema() throws IOException { + Schema arrowSchema = new Schema(asList( + new Field("id", FieldType.nullable(new ArrowType.Int(32, true)), null), + new Field("name", FieldType.nullable(new ArrowType.Utf8()), null), + new Field("active", FieldType.nullable(new ArrowType.Bool()), null), + new Field("score", FieldType.nullable( + new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), null))); + + MessageType parquetSchema = new SchemaConverter().fromArrow(arrowSchema).getParquetSchema(); + Path path = tempDir.toPath().resolve("mixed.parquet"); + OutputFile outputFile = new LocalOutputFile(path); + + try (VectorSchemaRoot batch = VectorSchemaRoot.create(arrowSchema, allocator)) { + IntVector idVector = (IntVector) batch.getVector("id"); + org.apache.arrow.vector.VarCharVector nameVector = + (org.apache.arrow.vector.VarCharVector) batch.getVector("name"); + org.apache.arrow.vector.BitVector activeVector = + (org.apache.arrow.vector.BitVector) batch.getVector("active"); + Float8Vector scoreVector = (Float8Vector) batch.getVector("score"); + + idVector.allocateNew(3); + nameVector.allocateNew(); + activeVector.allocateNew(3); + scoreVector.allocateNew(3); + + idVector.set(0, 1); + nameVector.set(0, "alice".getBytes(java.nio.charset.StandardCharsets.UTF_8)); + activeVector.set(0, 1); + scoreVector.set(0, 95.5); + + idVector.set(1, 2); + nameVector.set(1, "bob".getBytes(java.nio.charset.StandardCharsets.UTF_8)); + activeVector.set(1, 0); + scoreVector.setNull(1); + + idVector.setNull(2); + nameVector.setNull(2); + activeVector.setNull(2); + scoreVector.set(2, 77.0); + + batch.setRowCount(3); + + try (ArrowParquetWriter writer = new ArrowParquetWriter(outputFile, parquetSchema)) { + writer.writeBatch(batch); + } + } + + try (ParquetReader reader = + ParquetReader.builder(new GroupReadSupport(), new org.apache.hadoop.fs.Path(path.toUri())).build()) { + Group row0 = reader.read(); + assertThat(row0.getInteger("id", 0)).isEqualTo(1); + assertThat(row0.getBinary("name", 0).toStringUsingUTF8()).isEqualTo("alice"); + assertThat(row0.getBoolean("active", 0)).isTrue(); + assertThat(row0.getDouble("score", 0)).isEqualTo(95.5); + + Group row1 = reader.read(); + assertThat(row1.getInteger("id", 0)).isEqualTo(2); + assertThat(row1.getBinary("name", 0).toStringUsingUTF8()).isEqualTo("bob"); + assertThat(row1.getBoolean("active", 0)).isFalse(); + assertThat(row1.getFieldRepetitionCount("score")).isEqualTo(0); // null + + Group row2 = reader.read(); + assertThat(row2.getFieldRepetitionCount("id")).isEqualTo(0); // null + assertThat(row2.getFieldRepetitionCount("name")).isEqualTo(0); // null + assertThat(row2.getFieldRepetitionCount("active")).isEqualTo(0); // null + assertThat(row2.getDouble("score", 0)).isEqualTo(77.0); + + assertThat(reader.read()).isNull(); + } + } +} + +