diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index b7c88254c9..be9ac4f6fb 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -326,6 +326,7 @@ jobs: org.apache.comet.exec.CometShuffleSuite org.apache.comet.exec.CometShuffle4_0Suite org.apache.comet.exec.CometNativeColumnarToRowSuite + org.apache.comet.exec.CometDirectColumnarToRowSuite org.apache.comet.exec.CometNativeShuffleSuite org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite org.apache.comet.exec.CometShuffleEncryptionSuite @@ -344,6 +345,7 @@ jobs: org.apache.comet.exec.CometJoinSuite org.apache.spark.sql.comet.CometMapInBatchSuite org.apache.comet.CometNativeSuite + org.apache.comet.DirectColumnarToRowConverterSuite org.apache.comet.CometConfSuite org.apache.comet.CometPublicApiSuite org.apache.comet.QueryContextInternerSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 8210f91b7f..1dce1c8c88 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -142,6 +142,7 @@ jobs: org.apache.comet.exec.CometShuffleSuite org.apache.comet.exec.CometShuffle4_0Suite org.apache.comet.exec.CometNativeColumnarToRowSuite + org.apache.comet.exec.CometDirectColumnarToRowSuite org.apache.comet.exec.CometNativeShuffleSuite org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite org.apache.comet.exec.CometShuffleEncryptionSuite @@ -160,6 +161,7 @@ jobs: org.apache.comet.exec.CometJoinSuite org.apache.spark.sql.comet.CometMapInBatchSuite org.apache.comet.CometNativeSuite + org.apache.comet.DirectColumnarToRowConverterSuite org.apache.comet.CometConfSuite org.apache.comet.CometPublicApiSuite org.apache.comet.QueryContextInternerSuite diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index 5499052464..63fb3c3e2c 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -284,6 +284,21 @@ plan to fall back to Spark row-based execution — Comet removes its native oper mix of native and fallback operators joined by repeated conversions — which can be cheaper than paying the expensive conversions again and again. +### Experimental: Direct Columnar-to-Row Conversion + +When the JVM columnar-to-row operator is in use (`spark.comet.exec.columnarToRow.native.enabled=false`), +setting `spark.comet.exec.columnarToRow.direct.enabled=true` enables an experimental converter that writes +values straight from Arrow buffers into Spark's row format without allocating an object per value. This is +most beneficial for decimal-heavy schemas, where the default conversion allocates a `Decimal` object per +value (and considerably more for decimals with precision above 18); microbenchmarks show up to 2x faster +conversion and a large reduction in garbage creation for such schemas. Schemas containing data types the +converter does not support fall back to the default conversion automatically. + +Batches with fewer rows than `spark.comet.exec.columnarToRow.direct.minBatchSize` (default `128`) also fall +back to the default conversion, since the direct converter's per-batch setup does not pay off on very small +batches. This optimization is experimental: it only affects the operator's non-codegen paths (including +broadcast relation builds), and the default conversion remains enabled unless explicitly opted in. + ## Metrics Overhead Comet exposes rich native operator metrics for observability (see [Metrics](metrics.md)), but they are diff --git a/spark/src/main/java/org/apache/comet/DirectColumnarToRowConverter.java b/spark/src/main/java/org/apache/comet/DirectColumnarToRowConverter.java new file mode 100644 index 0000000000..f80de7aa76 --- /dev/null +++ b/spark/src/main/java/org/apache/comet/DirectColumnarToRowConverter.java @@ -0,0 +1,521 @@ +/* + * 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.comet; + +import org.apache.spark.sql.catalyst.expressions.UnsafeRow; +import org.apache.spark.sql.types.*; +import org.apache.spark.sql.vectorized.ColumnVector; +import org.apache.spark.sql.vectorized.ColumnarBatch; +import org.apache.spark.unsafe.Platform; +import org.apache.spark.unsafe.array.ByteArrayMethods; +import org.apache.spark.unsafe.types.UTF8String; + +import org.apache.comet.vector.CometPlainVector; +import org.apache.comet.vector.CometVector; + +/** + * Converts columnar batches to {@link UnsafeRow}s without per-value object allocation. + * + *
The conversion done by {@code rowIterator} plus {@code UnsafeProjection} routes every value + * through two virtual calls and allocates per object-typed value: a {@code Decimal} per compact + * decimal value and a {@code byte[]}/{@code BigInteger}/{@code BigDecimal} chain per decimal value + * with precision above 18. This converter instead resolves each column to a primitive type code + * once at construction and writes values straight from the Arrow buffers into a reused row buffer: + * compact decimals as unscaled longs via {@link CometVector#getLongDecimal(int)} and wide decimals + * as raw big-endian bytes via {@link CometVector#copyBinaryDecimal(int, byte[])}. + * + *
When every column is fixed-width (no strings, no decimals above precision 18), the whole batch + * is converted column-at-a-time in {@link #setBatch(ColumnarBatch)} into a single buffer with a + * constant row stride: one monomorphic loop per column with the null check hoisted out for + * null-free columns, mirroring the native converter's fixed-width fast path. In that mode {@link + * #convertRow(int)} only repoints the reused row. + * + *
The produced rows are byte-identical to {@code UnsafeProjection} output, which matters because + * {@link UnsafeRow} equality and hashing are byte-wise. The returned row and its backing buffer are + * reused across calls; callers must copy a row to retain it. + */ +public final class DirectColumnarToRowConverter { + + private static final int BOOLEAN = 0; + private static final int BYTE = 1; + private static final int SHORT = 2; + private static final int INT = 3; + private static final int LONG = 4; + private static final int FLOAT = 5; + private static final int DOUBLE = 6; + private static final int STRING = 7; + private static final int DECIMAL_COMPACT = 8; + private static final int DECIMAL_WIDE = 9; + + private final int numFields; + private final int[] typeCodes; + private final int[] precisions; + private final int[] scales; + private final int nullBitsetWidth; + private final int fixedSize; + private final boolean allFixedWidth; + + private final byte[] decimalBytes = new byte[16]; + private final UnsafeRow row; + private byte[] buffer = new byte[64]; + private int cursor; + + // Per-batch state + private ColumnVector[] columns; + private boolean[] hasNulls; + + // Fixed-width fast path state: the whole batch converted at a constant row stride. + private byte[] batchBuffer = new byte[0]; + private int batchNumRows; + + /** Returns true if every field of the schema has a type this converter supports. */ + public static boolean supportsSchema(StructType schema) { + for (StructField field : schema.fields()) { + if (typeCodeFor(field.dataType()) < 0) { + return false; + } + } + return true; + } + + private static int typeCodeFor(DataType dt) { + if (dt instanceof BooleanType) { + return BOOLEAN; + } else if (dt instanceof ByteType) { + return BYTE; + } else if (dt instanceof ShortType) { + return SHORT; + } else if (dt instanceof IntegerType || dt instanceof DateType) { + return INT; + } else if (dt instanceof LongType + || dt instanceof TimestampType + || dt instanceof TimestampNTZType) { + return LONG; + } else if (dt instanceof FloatType) { + return FLOAT; + } else if (dt instanceof DoubleType) { + return DOUBLE; + } else if (dt instanceof StringType) { + return STRING; + } else if (dt instanceof DecimalType) { + return ((DecimalType) dt).precision() <= Decimal.MAX_LONG_DIGITS() + ? DECIMAL_COMPACT + : DECIMAL_WIDE; + } else { + return -1; + } + } + + public DirectColumnarToRowConverter(StructType schema) { + StructField[] fields = schema.fields(); + numFields = fields.length; + typeCodes = new int[numFields]; + precisions = new int[numFields]; + scales = new int[numFields]; + for (int i = 0; i < numFields; i++) { + DataType dt = fields[i].dataType(); + int code = typeCodeFor(dt); + if (code < 0) { + throw new UnsupportedOperationException( + "DirectColumnarToRowConverter does not support data type: " + dt); + } + typeCodes[i] = code; + if (dt instanceof DecimalType) { + DecimalType d = (DecimalType) dt; + precisions[i] = d.precision(); + scales[i] = d.scale(); + } + } + nullBitsetWidth = UnsafeRow.calculateBitSetWidthInBytes(numFields); + long fixedSizeLong = nullBitsetWidth + (long) numFields * 8; + if (fixedSizeLong > ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH) { + throw new IllegalArgumentException("Schema is too wide for UnsafeRow conversion"); + } + fixedSize = (int) fixedSizeLong; + boolean fixed = true; + for (int i = 0; i < numFields; i++) { + if (typeCodes[i] == STRING || typeCodes[i] == DECIMAL_WIDE) { + fixed = false; + break; + } + } + allFixedWidth = fixed; + row = new UnsafeRow(numFields); + if (buffer.length < fixedSize) { + buffer = new byte[fixedSize]; + } + } + + /** Prepares the converter for a new batch. */ + public void setBatch(ColumnarBatch batch) { + if (batch.numCols() != numFields) { + throw new IllegalArgumentException( + "Column count mismatch: expected " + numFields + ", got " + batch.numCols()); + } + if (columns == null) { + columns = new ColumnVector[numFields]; + hasNulls = new boolean[numFields]; + } + for (int i = 0; i < numFields; i++) { + columns[i] = batch.column(i); + hasNulls[i] = columns[i].hasNull(); + } + if (allFixedWidth) { + convertBatchFixedWidth(batch.numRows()); + } + } + + /** + * Converts one row of the current batch. The returned row is reused across calls and valid until + * the next call ({@code setBatch} for fixed-width schemas). + */ + public UnsafeRow convertRow(int rowId) { + if (allFixedWidth) { + row.pointTo(batchBuffer, Platform.BYTE_ARRAY_OFFSET + (long) rowId * fixedSize, fixedSize); + return row; + } + // Zero the null bitset; fixed slots are always fully written below. + for (int i = 0; i < nullBitsetWidth; i += 8) { + Platform.putLong(buffer, Platform.BYTE_ARRAY_OFFSET + i, 0L); + } + cursor = fixedSize; + + for (int c = 0; c < numFields; c++) { + ColumnVector col = columns[c]; + long slot = Platform.BYTE_ARRAY_OFFSET + nullBitsetWidth + c * 8L; + boolean isNull = hasNulls[c] && col.isNullAt(rowId); + switch (typeCodes[c]) { + case BOOLEAN: + if (isNull) { + setNull(c, slot); + } else { + Platform.putLong(buffer, slot, 0L); + Platform.putBoolean(buffer, slot, col.getBoolean(rowId)); + } + break; + case BYTE: + if (isNull) { + setNull(c, slot); + } else { + Platform.putLong(buffer, slot, 0L); + Platform.putByte(buffer, slot, col.getByte(rowId)); + } + break; + case SHORT: + if (isNull) { + setNull(c, slot); + } else { + Platform.putLong(buffer, slot, 0L); + Platform.putShort(buffer, slot, col.getShort(rowId)); + } + break; + case INT: + if (isNull) { + setNull(c, slot); + } else { + Platform.putLong(buffer, slot, 0L); + Platform.putInt(buffer, slot, col.getInt(rowId)); + } + break; + case LONG: + if (isNull) { + setNull(c, slot); + } else { + Platform.putLong(buffer, slot, col.getLong(rowId)); + } + break; + case FLOAT: + if (isNull) { + setNull(c, slot); + } else { + Platform.putLong(buffer, slot, 0L); + Platform.putFloat(buffer, slot, col.getFloat(rowId)); + } + break; + case DOUBLE: + if (isNull) { + setNull(c, slot); + } else { + Platform.putDouble(buffer, slot, col.getDouble(rowId)); + } + break; + case STRING: + if (isNull) { + setNull(c, slot); + } else { + writeString(slot, col.getUTF8String(rowId)); + } + break; + case DECIMAL_COMPACT: + if (isNull) { + setNull(c, slot); + } else { + Platform.putLong(buffer, slot, compactDecimalValue(col, rowId, c)); + } + break; + case DECIMAL_WIDE: + writeWideDecimal(c, slot, col, rowId, isNull); + break; + default: + throw new IllegalStateException("Unknown type code: " + typeCodes[c]); + } + } + + row.pointTo(buffer, cursor); + return row; + } + + private void setNull(int ordinal, long slot) { + // Matches UnsafeRowWriter.setNullAt: set the bit and zero the fixed slot. + long wordOffset = Platform.BYTE_ARRAY_OFFSET + (ordinal >> 6) * 8L; + long word = Platform.getLong(buffer, wordOffset); + Platform.putLong(buffer, wordOffset, word | (1L << (ordinal & 63))); + Platform.putLong(buffer, slot, 0L); + } + + private void writeString(long slot, UTF8String value) { + int numBytes = value.numBytes(); + long roundedSize = ((long) numBytes + 7) & ~7L; + ensureCapacity((long) cursor + roundedSize); + if ((numBytes & 7) != 0) { + // Zero the last partial word so padding bytes are deterministic (buffer is reused). + Platform.putLong(buffer, Platform.BYTE_ARRAY_OFFSET + cursor + ((numBytes >> 3) << 3), 0L); + } + value.writeToMemory(buffer, Platform.BYTE_ARRAY_OFFSET + cursor); + Platform.putLong(buffer, slot, ((long) cursor << 32) | numBytes); + cursor += (int) roundedSize; + } + + private void writeWideDecimal( + int ordinal, long slot, ColumnVector col, int rowId, boolean isNull) { + // Matches UnsafeRowWriter.write(ordinal, Decimal, precision, scale) for precision > 18: + // 16 bytes are always reserved (and consumed) in the variable-length region, the minimal + // big-endian two's-complement bytes are written at the cursor, and for null values the null + // bit is set while the offset is still recorded with size 0. + ensureCapacity((long) cursor + 16); + Platform.putLong(buffer, Platform.BYTE_ARRAY_OFFSET + cursor, 0L); + Platform.putLong(buffer, Platform.BYTE_ARRAY_OFFSET + cursor + 8, 0L); + if (isNull) { + long wordOffset = Platform.BYTE_ARRAY_OFFSET + (ordinal >> 6) * 8L; + long word = Platform.getLong(buffer, wordOffset); + Platform.putLong(buffer, wordOffset, word | (1L << (ordinal & 63))); + Platform.putLong(buffer, slot, (long) cursor << 32); + } else { + byte[] be; + int start; + if (col instanceof CometPlainVector) { + be = ((CometVector) col).copyBinaryDecimal(rowId, decimalBytes); + // Trim to the minimal two's-complement form BigInteger.toByteArray would produce, + // so the row bytes match what UnsafeRowWriter writes. + byte sign = (be[0] & 0x80) != 0 ? (byte) 0xFF : (byte) 0x00; + start = 0; + while (start < 15 && be[start] == sign && ((be[start + 1] ^ sign) & 0x80) == 0) { + start++; + } + } else { + // Dictionary-encoded or other vector: fall back to the allocating accessor. + be = + col.getDecimal(rowId, precisions[ordinal], scales[ordinal]) + .toJavaBigDecimal() + .unscaledValue() + .toByteArray(); + start = 0; + } + int numBytes = be.length - start; + Platform.copyMemory( + be, + Platform.BYTE_ARRAY_OFFSET + start, + buffer, + Platform.BYTE_ARRAY_OFFSET + cursor, + numBytes); + Platform.putLong(buffer, slot, ((long) cursor << 32) | numBytes); + } + cursor += 16; + } + + private void convertBatchFixedWidth(int numRows) { + batchNumRows = numRows; + long totalSize = (long) fixedSize * numRows; + if (totalSize > ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH) { + throw new IllegalArgumentException("Batch too large for fixed-width conversion"); + } + if (batchBuffer.length < totalSize) { + batchBuffer = new byte[(int) totalSize]; + } + // Zero the null bitset of every row; fixed slots are always fully written below. + for (int r = 0; r < numRows; r++) { + long rowStart = Platform.BYTE_ARRAY_OFFSET + (long) r * fixedSize; + for (int w = 0; w < nullBitsetWidth; w += 8) { + Platform.putLong(batchBuffer, rowStart + w, 0L); + } + } + for (int c = 0; c < numFields; c++) { + writeColumnFixedWidth(c); + } + } + + /** + * Writes one column's values for all rows of the batch. Each type gets its own loop so the + * accessor call site stays monomorphic, and every slot is written as a single long whose byte + * layout matches what UnsafeRowWriter produces for that type (values are little-endian, so a + * narrow value zero-extended to a long occupies the same bytes as a partial write into a zeroed + * slot). + */ + private void writeColumnFixedWidth(int c) { + ColumnVector col = columns[c]; + boolean mayHaveNulls = hasNulls[c]; + int stride = fixedSize; + int n = batchNumRows; + long slotBase = Platform.BYTE_ARRAY_OFFSET + nullBitsetWidth + c * 8L; + long bitWordBase = Platform.BYTE_ARRAY_OFFSET + (c >> 6) * 8L; + long bitMask = 1L << (c & 63); + switch (typeCodes[c]) { + case BOOLEAN: + for (int r = 0; r < n; r++) { + long slot = slotBase + (long) r * stride; + if (mayHaveNulls && col.isNullAt(r)) { + setNullFixedWidth(r, stride, slot, bitWordBase, bitMask); + } else { + Platform.putLong(batchBuffer, slot, col.getBoolean(r) ? 1L : 0L); + } + } + break; + case BYTE: + for (int r = 0; r < n; r++) { + long slot = slotBase + (long) r * stride; + if (mayHaveNulls && col.isNullAt(r)) { + setNullFixedWidth(r, stride, slot, bitWordBase, bitMask); + } else { + Platform.putLong(batchBuffer, slot, col.getByte(r) & 0xFFL); + } + } + break; + case SHORT: + for (int r = 0; r < n; r++) { + long slot = slotBase + (long) r * stride; + if (mayHaveNulls && col.isNullAt(r)) { + setNullFixedWidth(r, stride, slot, bitWordBase, bitMask); + } else { + Platform.putLong(batchBuffer, slot, col.getShort(r) & 0xFFFFL); + } + } + break; + case INT: + for (int r = 0; r < n; r++) { + long slot = slotBase + (long) r * stride; + if (mayHaveNulls && col.isNullAt(r)) { + setNullFixedWidth(r, stride, slot, bitWordBase, bitMask); + } else { + Platform.putLong(batchBuffer, slot, col.getInt(r) & 0xFFFFFFFFL); + } + } + break; + case LONG: + for (int r = 0; r < n; r++) { + long slot = slotBase + (long) r * stride; + if (mayHaveNulls && col.isNullAt(r)) { + setNullFixedWidth(r, stride, slot, bitWordBase, bitMask); + } else { + Platform.putLong(batchBuffer, slot, col.getLong(r)); + } + } + break; + case FLOAT: + for (int r = 0; r < n; r++) { + long slot = slotBase + (long) r * stride; + if (mayHaveNulls && col.isNullAt(r)) { + setNullFixedWidth(r, stride, slot, bitWordBase, bitMask); + } else { + Platform.putLong( + batchBuffer, slot, Float.floatToRawIntBits(col.getFloat(r)) & 0xFFFFFFFFL); + } + } + break; + case DOUBLE: + for (int r = 0; r < n; r++) { + long slot = slotBase + (long) r * stride; + if (mayHaveNulls && col.isNullAt(r)) { + setNullFixedWidth(r, stride, slot, bitWordBase, bitMask); + } else { + Platform.putLong(batchBuffer, slot, Double.doubleToRawLongBits(col.getDouble(r))); + } + } + break; + case DECIMAL_COMPACT: + if (col instanceof CometVector) { + CometVector cometCol = (CometVector) col; + for (int r = 0; r < n; r++) { + long slot = slotBase + (long) r * stride; + if (mayHaveNulls && col.isNullAt(r)) { + setNullFixedWidth(r, stride, slot, bitWordBase, bitMask); + } else { + Platform.putLong(batchBuffer, slot, cometCol.getLongDecimal(r)); + } + } + } else { + for (int r = 0; r < n; r++) { + long slot = slotBase + (long) r * stride; + if (mayHaveNulls && col.isNullAt(r)) { + setNullFixedWidth(r, stride, slot, bitWordBase, bitMask); + } else { + Platform.putLong(batchBuffer, slot, compactDecimalValue(col, r, c)); + } + } + } + break; + default: + throw new IllegalStateException( + "Type code not supported by fixed-width path: " + typeCodes[c]); + } + } + + /** + * Reads a compact decimal's unscaled long. Comet vectors expose it allocation-free; other vector + * types (e.g. ConstantColumnVector) go through the allocating accessor. + */ + private long compactDecimalValue(ColumnVector col, int rowId, int ordinal) { + if (col instanceof CometVector) { + return ((CometVector) col).getLongDecimal(rowId); + } + return col.getDecimal(rowId, precisions[ordinal], scales[ordinal]).toUnscaledLong(); + } + + private void setNullFixedWidth(int rowId, int stride, long slot, long bitWordBase, long bitMask) { + long wordAddr = bitWordBase + (long) rowId * stride; + Platform.putLong(batchBuffer, wordAddr, Platform.getLong(batchBuffer, wordAddr) | bitMask); + Platform.putLong(batchBuffer, slot, 0L); + } + + private void ensureCapacity(long needed) { + if (needed > ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH) { + throw new IllegalArgumentException("Row too large for UnsafeRow conversion"); + } + if (needed > buffer.length) { + int newSize = + (int) + Math.min( + Math.max(needed, (long) buffer.length * 2), + ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH); + byte[] newBuffer = new byte[newSize]; + System.arraycopy(buffer, 0, newBuffer, 0, cursor); + buffer = newBuffer; + } + } +} diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index d8fe5b6989..6dd8bd0bfc 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -269,6 +269,32 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_DIRECT_COLUMNAR_TO_ROW_ENABLED: ConfigEntry[Boolean] = + conf(s"$COMET_EXEC_CONFIG_PREFIX.columnarToRow.direct.enabled") + .category(CATEGORY_EXEC) + .doc( + "Experimental: Whether to use the direct columnar to row converter in the JVM " + + "columnar to row operator's non-codegen paths (including broadcast relation " + + "builds). The direct converter writes values straight from Arrow buffers into " + + "UnsafeRow format without per-value object allocation, which is significantly " + + "faster for decimal-heavy schemas. Schemas with unsupported data types fall back " + + "to the default conversion. Only applies when the JVM columnar to row operator is " + + "used (spark.comet.exec.columnarToRow.native.enabled=false).") + .booleanConf + .createWithDefault(false) + + val COMET_DIRECT_COLUMNAR_TO_ROW_MIN_BATCH_SIZE: ConfigEntry[Int] = + conf(s"$COMET_EXEC_CONFIG_PREFIX.columnarToRow.direct.minBatchSize") + .category(CATEGORY_EXEC) + .doc( + "Experimental: Batches with fewer rows than this use the default conversion even " + + s"when $COMET_EXEC_CONFIG_PREFIX.columnarToRow.direct.enabled is true, since the " + + "direct converter's per-batch setup does not amortize on very small batches. Only " + + "applies when the direct columnar to row converter is enabled.") + .intConf + .checkValue(_ >= 0, "Must be >= 0.") + .createWithDefault(128) + val COMET_EXEC_SORT_MERGE_JOIN_WITH_JOIN_FILTER_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.exec.sortMergeJoinWithJoinFilter.enabled") .category(CATEGORY_ENABLE_EXEC) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala index 2fe870ed06..d52ce10ac6 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala @@ -45,6 +45,8 @@ import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.spark.util.{SparkFatalException, Utils} import org.apache.spark.util.io.ChunkedByteBuffer +import org.apache.comet.{CometConf, DirectColumnarToRowConverter} + /** * Copied from Spark `ColumnarToRowExec`. Comet needs the fix for SPARK-50235 but cannot wait for * the fix to be released in Spark versions. We copy the implementation here to apply the fix. @@ -70,19 +72,31 @@ case class CometColumnarToRowExec(child: SparkPlan) "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"), "numInputBatches" -> SQLMetrics.createMetric(sparkContext, "number of input batches")) + private def useDirectConverter: Boolean = + CometConf.COMET_DIRECT_COLUMNAR_TO_ROW_ENABLED.get(conf) && + DirectColumnarToRowConverter.supportsSchema( + StructType(output.map(a => StructField(a.name, a.dataType, a.nullable)))) + + private def directConverterMinBatchSize: Int = + CometConf.COMET_DIRECT_COLUMNAR_TO_ROW_MIN_BATCH_SIZE.get(conf) + override def doExecute(): RDD[InternalRow] = { val numOutputRows = longMetric("numOutputRows") val numInputBatches = longMetric("numInputBatches") // This avoids calling `output` in the RDD closure, so that we don't need to include the entire // plan (this) in the closure. val localOutput = this.output + val direct = useDirectConverter + val minBatchSize = directConverterMinBatchSize child.executeColumnar().mapPartitionsInternal { batches => - val toUnsafe = UnsafeProjection.create(localOutput, localOutput) - batches.flatMap { batch => - numInputBatches += 1 - numOutputRows += batch.numRows() - batch.rowIterator().asScala.map(toUnsafe) - } + CometColumnarToRowExec + .convertBatches( + batches, + localOutput, + direct, + minBatchSize, + numInputBatches, + numOutputRows) } } @@ -113,14 +127,15 @@ case class CometColumnarToRowExec(child: SparkPlan) val localOutput = this.output val broadcastColumnar = child.executeBroadcast() val serializedBatches = broadcastColumnar.value.asInstanceOf[Array[ChunkedByteBuffer]] - val toUnsafe = UnsafeProjection.create(localOutput, localOutput) - val rows = serializedBatches.iterator + val batches = serializedBatches.iterator .flatMap(CometUtils.decodeBatches(_, this.getClass.getSimpleName)) - .flatMap { batch => - numInputBatches += 1 - numOutputRows += batch.numRows() - batch.rowIterator().asScala.map(toUnsafe) - } + val rows = CometColumnarToRowExec.convertBatches( + batches, + localOutput, + useDirectConverter, + directConverterMinBatchSize, + numInputBatches, + numOutputRows) val mode = cometBroadcastExchange.get.mode val relation = mode.transform(rows, Some(numOutputRows.value)) @@ -303,3 +318,54 @@ case class CometColumnarToRowExec(child: SparkPlan) override protected def withNewChildInternal(newChild: SparkPlan): CometColumnarToRowExec = copy(child = newChild) } + +object CometColumnarToRowExec { + + /** + * Converts columnar batches to rows, either through the allocation-free + * [[DirectColumnarToRowConverter]] or the default `rowIterator` plus `UnsafeProjection` path. + * When the direct converter is enabled, batches smaller than `minBatchSize` still use the + * default path since the direct converter's per-batch setup does not amortize on very small + * batches. Both paths reuse the returned row across calls; consumers must copy rows they + * retain. + */ + private[comet] def convertBatches( + batches: Iterator[ColumnarBatch], + output: Seq[Attribute], + useDirectConverter: Boolean, + minBatchSize: Int, + numInputBatches: SQLMetric, + numOutputRows: SQLMetric): Iterator[InternalRow] = { + if (useDirectConverter) { + val schema = StructType(output.map(a => StructField(a.name, a.dataType, a.nullable))) + val converter = new DirectColumnarToRowConverter(schema) + lazy val toUnsafe = UnsafeProjection.create(output, output) + batches.flatMap { batch => + numInputBatches += 1 + numOutputRows += batch.numRows() + if (batch.numRows() >= minBatchSize) { + converter.setBatch(batch) + val numRows = batch.numRows() + new Iterator[InternalRow] { + private var i = 0 + override def hasNext: Boolean = i < numRows + override def next(): InternalRow = { + val row = converter.convertRow(i) + i += 1 + row + } + } + } else { + batch.rowIterator().asScala.map(toUnsafe) + } + } + } else { + val toUnsafe = UnsafeProjection.create(output, output) + batches.flatMap { batch => + numInputBatches += 1 + numOutputRows += batch.numRows() + batch.rowIterator().asScala.map(toUnsafe) + } + } + } +} diff --git a/spark/src/test/scala/org/apache/comet/DirectColumnarToRowConverterSuite.scala b/spark/src/test/scala/org/apache/comet/DirectColumnarToRowConverterSuite.scala new file mode 100644 index 0000000000..9c972b2f9f --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/DirectColumnarToRowConverterSuite.scala @@ -0,0 +1,230 @@ +/* + * 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.comet + +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeProjection, UnsafeRow} +import org.apache.spark.sql.comet.execution.arrow.CometArrowConverters +import org.apache.spark.sql.execution.vectorized.ConstantColumnVector +import org.apache.spark.sql.types._ +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} +import org.apache.spark.unsafe.types.UTF8String + +/** + * Verifies that [[DirectColumnarToRowConverter]] produces rows byte-identical to + * `UnsafeProjection`, which matters because `UnsafeRow` equality and hashing are byte-wise. + */ +class DirectColumnarToRowConverterSuite extends AnyFunSuite { + + private def assertMatchesUnsafeProjection( + schema: StructType, + rows: Seq[InternalRow], + batchSize: Int): Unit = { + val batches = CometArrowConverters + .rowToArrowBatchIter(rows.iterator, schema, batchSize, "UTC", CometArrowAllocator) + .toArray + try { + val proj = UnsafeProjection.create(schema.fields.map(_.dataType)) + val converter = new DirectColumnarToRowConverter(schema) + var rowIdx = 0 + for (batch <- batches) { + converter.setBatch(batch) + val it = batch.rowIterator() + var i = 0 + while (it.hasNext) { + val expected = proj(it.next()) + val actual = converter.convertRow(i) + assert( + expected.getBytes.toSeq == actual.getBytes.toSeq, + s"row $rowIdx differs:\n expected ${expected.getBytes.toSeq}\n actual " + + s"${actual.getBytes.toSeq}") + i += 1 + rowIdx += 1 + } + } + } finally { + batches.foreach(_.close()) + } + } + + test("all supported types with nulls match UnsafeProjection bytes") { + val schema = new StructType() + .add("bool", BooleanType) + .add("byte", ByteType) + .add("short", ShortType) + .add("int", IntegerType) + .add("long", LongType) + .add("float", FloatType) + .add("double", DoubleType) + .add("date", DateType) + .add("ts", TimestampType) + .add("ts_ntz", TimestampNTZType) + .add("str", StringType) + .add("dec_compact", DecimalType(12, 2)) + .add("dec_wide", DecimalType(38, 10)) + + val rows = (0 until 1000).map { i => + def nullEvery(k: Int, v: Any): Any = if (i % k == 0) null else v + val sign = if (i % 2 == 0) 1L else -1L + new GenericInternalRow( + Array[Any]( + nullEvery(3, i % 2 == 0), + nullEvery(4, ((i % 256) - 128).toByte), + nullEvery(5, (i - 500).toShort), + nullEvery(6, i * sign.toInt), + nullEvery(7, i.toLong * sign * 1000003L), + nullEvery(8, if (i % 50 == 1) Float.NaN else i.toFloat * sign), + nullEvery(9, if (i % 50 == 2) Double.NaN else i.toDouble * sign), + nullEvery(10, 8000 + i % 2500), + nullEvery(11, i.toLong * 1000000L), + nullEvery(13, i.toLong * -1000000L), + nullEvery(12, UTF8String.fromString(if (i % 13 == 0) "" else s"value_$i")), + nullEvery(3, Decimal.createUnsafe(i * sign * 97, 12, 2)), + nullEvery( + 4, { + val magnitude = new java.math.BigInteger(s"${i + 1}" * 3) + val unscaled = if (sign < 0) magnitude.negate() else magnitude + Decimal(new java.math.BigDecimal(unscaled, 10)) + }))) + } + + assertMatchesUnsafeProjection(schema, rows, batchSize = 64) + } + + test("all-fixed-width schema takes the columnar fast path and matches bytes") { + val schema = new StructType() + .add("bool", BooleanType) + .add("byte", ByteType) + .add("short", ShortType) + .add("int", IntegerType) + .add("long", LongType) + .add("float", FloatType) + .add("double", DoubleType) + .add("date", DateType) + .add("ts", TimestampType) + .add("ts_ntz", TimestampNTZType) + .add("dec_compact", DecimalType(12, 2)) + + val rows = (0 until 1000).map { i => + def nullEvery(k: Int, v: Any): Any = if (i % k == 0) null else v + val sign = if (i % 2 == 0) 1L else -1L + new GenericInternalRow( + Array[Any]( + nullEvery(3, i % 2 == 0), + nullEvery(4, ((i % 256) - 128).toByte), + nullEvery(5, (i - 500).toShort), + nullEvery(6, i * sign.toInt), + nullEvery(7, i.toLong * sign * 1000003L), + nullEvery(8, if (i % 50 == 1) Float.NaN else i.toFloat * sign), + nullEvery(9, if (i % 50 == 2) Double.NaN else i.toDouble * sign), + nullEvery(10, 8000 + i % 2500), + nullEvery(11, i.toLong * 1000000L), + nullEvery(13, i.toLong * -1000000L), + nullEvery(3, Decimal.createUnsafe(i * sign * 97, 12, 2)))) + } + + assertMatchesUnsafeProjection(schema, rows, batchSize = 64) + } + + test("multi-word null bitset matches UnsafeProjection bytes") { + val numCols = 70 + val schema = + (0 until numCols).foldLeft(new StructType())((s, i) => s.add(s"c$i", LongType)) + val rows = (0 until 200).map { i => + new GenericInternalRow((0 until numCols).map { c => + if ((i + c) % 3 == 0) null else (i.toLong * 31 + c): Any + }.toArray) + } + assertMatchesUnsafeProjection(schema, rows, batchSize = 33) + } + + test("wide decimal boundary values match UnsafeProjection bytes") { + val schema = new StructType().add("d", DecimalType(38, 0)) + val big = new java.math.BigInteger("99999999999999999999999999999999999999") + val values = Seq( + java.math.BigInteger.ZERO, + java.math.BigInteger.ONE, + java.math.BigInteger.ONE.negate(), + java.math.BigInteger.valueOf(Long.MaxValue), + java.math.BigInteger.valueOf(Long.MinValue), + java.math.BigInteger.valueOf(127), + java.math.BigInteger.valueOf(128), + java.math.BigInteger.valueOf(-128), + java.math.BigInteger.valueOf(-129), + big, + big.negate()) + val rows = values.map { v => + new GenericInternalRow(Array[Any](Decimal(new java.math.BigDecimal(v, 0)))) + } + assertMatchesUnsafeProjection(schema, rows, batchSize = 4) + } + + test("noncanonical NaN payloads match UnsafeProjection bytes") { + val floatValue = java.lang.Float.intBitsToFloat(0x7fc12345) + val doubleValue = java.lang.Double.longBitsToDouble(0x7ff8000000000001L) + + for (includeString <- Seq(false, true)) { + val fixedSchema = new StructType().add("f", FloatType).add("d", DoubleType) + val schema = if (includeString) fixedSchema.add("s", StringType) else fixedSchema + val floatCol = new ConstantColumnVector(1, FloatType) + floatCol.setFloat(floatValue) + val doubleCol = new ConstantColumnVector(1, DoubleType) + doubleCol.setDouble(doubleValue) + val columns = if (includeString) { + val stringCol = new ConstantColumnVector(1, StringType) + stringCol.setUtf8String(UTF8String.fromString("general path")) + Array[ColumnVector](floatCol, doubleCol, stringCol) + } else { + Array[ColumnVector](floatCol, doubleCol) + } + val batch = new ColumnarBatch(columns, 1) + + try { + val expected = UnsafeProjection + .create(schema.fields.map(_.dataType))(batch.getRow(0)) + .getBytes + val converter = new DirectColumnarToRowConverter(schema) + converter.setBatch(batch) + assert(expected.sameElements(converter.convertRow(0).getBytes)) + } finally { + batch.close() + } + } + } + + test("oversized fixed-width batch is rejected before allocation") { + val schema = new StructType().add("value", LongType) + val rowSize = UnsafeRow.calculateBitSetWidthInBytes(1) + 8 + val numRows = + org.apache.spark.unsafe.array.ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH / rowSize + 1 + val column = new ConstantColumnVector(numRows, LongType) + val batch = new ColumnarBatch(Array[ColumnVector](column), numRows) + + try { + val converter = new DirectColumnarToRowConverter(schema) + val error = intercept[IllegalArgumentException](converter.setBatch(batch)) + assert(error.getMessage.contains("Batch too large")) + } finally { + batch.close() + } + } +} diff --git a/spark/src/test/scala/org/apache/comet/exec/CometDirectColumnarToRowSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometDirectColumnarToRowSuite.scala new file mode 100644 index 0000000000..1c5ff110f9 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometDirectColumnarToRowSuite.scala @@ -0,0 +1,103 @@ +/* + * 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.comet.exec + +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.comet.CometColumnarToRowExec +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.CometConf + +/** + * End-to-end tests for `spark.comet.exec.columnarToRow.direct.enabled`, which routes + * `CometColumnarToRowExec`'s non-codegen conversion through `DirectColumnarToRowConverter`. + * Whole-stage codegen is disabled so the operator's `doExecute` conversion path runs. + */ +class CometDirectColumnarToRowSuite extends CometTestBase { + + private def withDirectConverter(minBatchSize: Int = 0)(f: => Unit): Unit = { + withSQLConf( + CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "false", + CometConf.COMET_DIRECT_COLUMNAR_TO_ROW_ENABLED.key -> "true", + CometConf.COMET_DIRECT_COLUMNAR_TO_ROW_MIN_BATCH_SIZE.key -> minBatchSize.toString, + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false")(f) + } + + private def checkQueryUsesJvmColumnarToRow(query: String): Unit = { + val df = sql(query) + val c2r = df.queryExecution.executedPlan.collect { case c: CometColumnarToRowExec => c } + assert(c2r.nonEmpty, s"expected CometColumnarToRowExec in plan:\n${df.queryExecution}") + checkSparkAnswer(df) + } + + test("direct converter: mixed types with nulls") { + withDirectConverter() { + withParquetTable((0 until 1000).map(i => (i.toLong, i.toString, i % 5)), "tbl") { + checkQueryUsesJvmColumnarToRow(""" + | SELECT + | _1, + | _2, + | CAST(_1 AS decimal(12,2)) AS dec_compact, + | CAST(_1 AS decimal(38,10)) AS dec_wide, + | DATE_ADD(DATE'2020-01-01', _3) AS dt, + | CAST(_1 AS double) AS dbl, + | CASE WHEN _3 = 0 THEN NULL ELSE _1 END AS maybe_null + | FROM tbl + |""".stripMargin) + } + } + } + + test("direct converter: all-fixed-width schema takes the columnar fast path") { + withDirectConverter() { + withParquetTable((0 until 1000).map(i => (i.toLong, i, i.toDouble)), "tbl") { + checkQueryUsesJvmColumnarToRow(""" + | SELECT + | _1, + | _2, + | _3, + | CAST(_1 AS decimal(10,2)) AS dec_compact, + | CASE WHEN _2 % 3 = 0 THEN NULL ELSE _2 END AS maybe_null + | FROM tbl + |""".stripMargin) + } + } + } + + test("batches below minBatchSize fall back to default conversion") { + withDirectConverter(minBatchSize = Int.MaxValue) { + withParquetTable((0 until 1000).map(i => (i.toLong, i.toString)), "tbl") { + // Every batch is below the threshold, so this exercises the per-batch fallback while + // the direct converter is enabled. + checkQueryUsesJvmColumnarToRow("SELECT _1, _2, CAST(_1 AS decimal(12,2)) AS dec FROM tbl") + } + } + } + + test("unsupported schema falls back to default conversion") { + withDirectConverter() { + withParquetTable((0 until 100).map(i => (i, i.toString)), "tbl") { + // BinaryType is not supported by DirectColumnarToRowConverter, so this exercises the + // per-plan fallback to the UnsafeProjection path. + checkQueryUsesJvmColumnarToRow("SELECT _1, CAST(_2 AS binary) FROM tbl") + } + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometC2RIsolatedBench.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometC2RIsolatedBench.scala index ce2526406e..2de899f5c6 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometC2RIsolatedBench.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometC2RIsolatedBench.scala @@ -27,41 +27,165 @@ import org.apache.spark.sql.types._ import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.spark.unsafe.types.UTF8String -import org.apache.comet.NativeColumnarToRowConverter +import org.apache.comet.{DirectColumnarToRowConverter, NativeColumnarToRowConverter} /** * Isolated columnar-to-row microbenchmark that excludes the parquet scan entirely. Batches are * built once in memory; each case converts them to rows and consumes the values so that the * conversion itself is what gets measured. + * + * The schema scenarios exist because the implementations have very different per-value costs by + * data type. The JVM path allocates a Decimal object per non-null decimal value (and a + * BigInteger/BigDecimal chain for precision > 18), while the native path writes the unscaled + * value directly. The native path also has a column-wise fast path that only engages when every + * column in the schema is fixed-width. */ object CometC2RIsolatedBench { private val totalRows = 1024 * 1024 - private def makeBatches(schema: StructType, batchSize: Int): Array[ColumnarBatch] = { - val rows = (0 until totalRows).iterator.map { i => - InternalRow(i.toLong, i, i.toDouble, UTF8String.fromString(s"value_$i")) - } + private case class Scenario( + name: String, + schema: StructType, + makeRow: Int => InternalRow, + consume: InternalRow => Long) + + private val mixedPrimitives = Scenario( + "long, int, double, string", + new StructType() + .add("a", LongType) + .add("b", IntegerType) + .add("c", DoubleType) + .add("d", StringType), + i => InternalRow(i.toLong, i, i.toDouble, UTF8String.fromString(s"value_$i")), + // All three converters return UnsafeRow, so reading a raw fixed slot avoids allocating a + // UTF8String in the benchmark sink. + row => row.getLong(0) + row.getLong(3)) + + // The columns a TPC-H lineitem aggregation reaches C2R with: four compact decimals, a date, + // and two short flag strings. The strings force the native general (row-at-a-time) path. + private val tpchDecimalsWithStrings = Scenario( + "4 x decimal(12,2), date, 2 x string (tpch-like)", + new StructType() + .add("l_quantity", DecimalType(12, 2)) + .add("l_extendedprice", DecimalType(12, 2)) + .add("l_discount", DecimalType(12, 2)) + .add("l_tax", DecimalType(12, 2)) + .add("l_shipdate", DateType) + .add("l_returnflag", StringType) + .add("l_linestatus", StringType), + i => + InternalRow( + Decimal.createUnsafe(i % 5000000, 12, 2), + Decimal.createUnsafe(i % 9000000, 12, 2), + Decimal.createUnsafe(i % 10, 12, 2), + Decimal.createUnsafe(i % 8, 12, 2), + 8000 + i % 2500, + UTF8String.fromString(if (i % 2 == 0) "A" else "R"), + UTF8String.fromString(if (i % 3 == 0) "F" else "O")), + // A compact decimal is stored directly as an unscaled long in UnsafeRow. + row => row.getLong(0) + row.getInt(4)) + + // Same decimals without any var-length column, so the native column-wise fast path engages. + private val decimalsAllFixedWidth = Scenario( + "4 x decimal(12,2), date, long (all fixed-width)", + new StructType() + .add("a", DecimalType(12, 2)) + .add("b", DecimalType(12, 2)) + .add("c", DecimalType(12, 2)) + .add("d", DecimalType(12, 2)) + .add("e", DateType) + .add("f", LongType), + i => + InternalRow( + Decimal.createUnsafe(i % 5000000, 12, 2), + Decimal.createUnsafe(i % 9000000, 12, 2), + Decimal.createUnsafe(i % 10, 12, 2), + Decimal.createUnsafe(i % 8, 12, 2), + 8000 + i % 2500, + i.toLong), + row => row.getLong(0) + row.getLong(5)) + + // Precision > 18: the JVM accessor allocates byte[] -> BigInteger -> java BigDecimal -> + // scala BigDecimal -> Decimal per value; the native side writes 16 bytes. + private val highPrecisionDecimals = Scenario( + "2 x decimal(38,10), long (high precision)", + new StructType() + .add("a", DecimalType(38, 10)) + .add("b", DecimalType(38, 10)) + .add("c", LongType), + i => + InternalRow( + Decimal(new java.math.BigDecimal(java.math.BigInteger.valueOf(i.toLong * 1000003), 10)), + Decimal(new java.math.BigDecimal(java.math.BigInteger.valueOf(i.toLong * 999983), 10)), + i.toLong), + // Consume only the cheap column: reading a decimal(38,10) back out of the UnsafeRow + // allocates on both sides and would mask the conversion cost being measured. + row => row.getLong(2)) + + private val widePrimitives = Scenario( + "16 x long (wide primitives)", + (0 until 16).foldLeft(new StructType())((s, i) => s.add(s"c$i", LongType)), + i => InternalRow((0 until 16).map(c => i.toLong + c): _*), + row => row.getLong(0) + row.getLong(15)) + + private def makeBatches(scenario: Scenario, batchSize: Int): Array[ColumnarBatch] = { + val rows = (0 until totalRows).iterator.map(scenario.makeRow) CometArrowConverters - .rowToArrowBatchIter(rows, schema, batchSize, "UTC", org.apache.comet.CometArrowAllocator) + .rowToArrowBatchIter( + rows, + scenario.schema, + batchSize, + "UTC", + org.apache.comet.CometArrowAllocator) .toArray } - private def runForBatchSize(schema: StructType, batchSize: Int): Unit = { - val batches = makeBatches(schema, batchSize) + /** Bytes allocated on the JVM heap by the current thread while running `body`. */ + private def measureAllocatedBytes(body: => Unit): Long = { + val bean = java.lang.management.ManagementFactory.getThreadMXBean + .asInstanceOf[com.sun.management.ThreadMXBean] + val tid = Thread.currentThread().getId + val before = bean.getThreadAllocatedBytes(tid) + body + bean.getThreadAllocatedBytes(tid) - before + } + + private def runForBatchSize(scenario: Scenario, batchSize: Int): Unit = { + val batches = makeBatches(scenario, batchSize) val benchmark = - new Benchmark(s"Isolated C2R (no scan), batchSize=$batchSize", totalRows.toLong) + new Benchmark( + s"Isolated C2R (no scan), ${scenario.name}, batchSize=$batchSize", + totalRows.toLong) benchmark.addCase("JVM rowIterator + UnsafeProjection") { _ => - val proj = UnsafeProjection.create(schema.fields.map(_.dataType)) + val proj = UnsafeProjection.create(scenario.schema.fields.map(_.dataType)) var sink = 0L var b = 0 while (b < batches.length) { val it = batches(b).rowIterator() while (it.hasNext) { val u = proj(it.next()) - sink += u.getLong(0) + u.getUTF8String(3).numBytes() + sink += scenario.consume(u) + } + b += 1 + } + if (sink == Long.MinValue) println(sink) + } + + benchmark.addCase("JVM direct converter (prototype)") { _ => + val converter = new DirectColumnarToRowConverter(scenario.schema) + var sink = 0L + var b = 0 + while (b < batches.length) { + val batch = batches(b) + converter.setBatch(batch) + val numRows = batch.numRows() + var i = 0 + while (i < numRows) { + sink += scenario.consume(converter.convertRow(i)) + i += 1 } b += 1 } @@ -69,15 +193,14 @@ object CometC2RIsolatedBench { } benchmark.addCase("Native converter") { _ => - val converter = new NativeColumnarToRowConverter(schema, batchSize) + val converter = new NativeColumnarToRowConverter(scenario.schema, batchSize) var sink = 0L try { var b = 0 while (b < batches.length) { val it = converter.convert(batches(b)) while (it.hasNext) { - val u = it.next() - sink += u.getLong(0) + u.getUTF8String(3).numBytes() + sink += scenario.consume(it.next()) } b += 1 } @@ -89,18 +212,68 @@ object CometC2RIsolatedBench { benchmark.run() + // GC pressure comparison: the JVM path allocates per object-typed value (Decimal, + // UTF8String) but reuses its output row buffer, while the native path allocates a + // byte[] + UnsafeRow per row for the defensive copy in NativeRowIterator. + val jvmAlloc = measureAllocatedBytes { + val proj = UnsafeProjection.create(scenario.schema.fields.map(_.dataType)) + var sink = 0L + for (batch <- batches) { + val it = batch.rowIterator() + while (it.hasNext) { + sink += scenario.consume(proj(it.next())) + } + } + if (sink == Long.MinValue) println(sink) + } + val directAlloc = measureAllocatedBytes { + val converter = new DirectColumnarToRowConverter(scenario.schema) + var sink = 0L + for (batch <- batches) { + converter.setBatch(batch) + val numRows = batch.numRows() + var i = 0 + while (i < numRows) { + sink += scenario.consume(converter.convertRow(i)) + i += 1 + } + } + if (sink == Long.MinValue) println(sink) + } + val nativeAlloc = measureAllocatedBytes { + val converter = new NativeColumnarToRowConverter(scenario.schema, batchSize) + var sink = 0L + try { + for (batch <- batches) { + val it = converter.convert(batch) + while (it.hasNext) { + sink += scenario.consume(it.next()) + } + } + } finally { + converter.close() + } + if (sink == Long.MinValue) println(sink) + } + println( + f"JVM heap allocation per row: JVM path ${jvmAlloc.toDouble / totalRows}%.1f bytes, " + + f"direct path ${directAlloc.toDouble / totalRows}%.1f bytes, " + + f"native path ${nativeAlloc.toDouble / totalRows}%.1f bytes%n") + batches.foreach(_.close()) } def main(args: Array[String]): Unit = { - val schema = new StructType() - .add("a", LongType) - .add("b", IntegerType) - .add("c", DoubleType) - .add("d", StringType) + val scenarios = + Seq( + mixedPrimitives, + tpchDecimalsWithStrings, + decimalsAllFixedWidth, + highPrecisionDecimals, + widePrimitives) - runForBatchSize(schema, 8192) - runForBatchSize(schema, 512) - runForBatchSize(schema, 32) + for (scenario <- scenarios; batchSize <- Seq(8192, 512, 32)) { + runForBatchSize(scenario, batchSize) + } } }