Skip to content

[SPARK-23188][SQL] Make vectorized columar reader batch size configurable #20361

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
jiangxb1987 wants to merge 3 commits into from
Closed

[SPARK-23188][SQL] Make vectorized columar reader batch size configurable #20361

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
927c6b4
make vector batch size configurable.
jiangxb1987 Jan 23, 2018
38debd7
fix
jiangxb1987 Jan 23, 2018
5ad935f
remove constructer VectorizedParquetRecordReader(useOffHeap).
jiangxb1987 Feb 1, 2018
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -377,6 +377,12 @@ object SQLConf {
.booleanConf
.createWithDefault(true)

val PARQUET_VECTORIZED_READER_BATCH_SIZE = buildConf("spark.sql.parquet.batchSize")

@cloud-fan cloud-fan Jan 30, 2018

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd prefer spark.sql.parquet.columnarReaderBatchSize to be more clear.

@gatorsmile gatorsmile Jan 30, 2018

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Still a question. Is that possible to use the estimated memory size instead of the number of rows?

@cloud-fan cloud-fan Jan 30, 2018

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd say it's very hard. If we need to satisfy a sizeInBytes limitation, we would need to load data record by record, and stop loading if we hit the limitation. But for performance reasons, we wanna load the data with batch, which needs to know the batch size ahead.

.doc("The number of rows to include in a parquet vectorized reader batch. The number should " +
"be carefully chosen to minimize overhead and avoid OOMs in reading data.")
.intConf
.createWithDefault(4096)

val ORC_COMPRESSION = buildConf("spark.sql.orc.compression.codec")
.doc("Sets the compression codec used when writing ORC files. If either `compression` or " +
"`orc.compress` is specified in the table-specific options/properties, the precedence " +
Expand All @@ -400,6 +406,12 @@ object SQLConf {
.booleanConf
.createWithDefault(true)

val ORC_VECTORIZED_READER_BATCH_SIZE = buildConf("spark.sql.orc.batchSize")

@cloud-fan cloud-fan Jan 30, 2018

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ditto

.doc("The number of rows to include in a orc vectorized reader batch. The number should " +
"be carefully chosen to minimize overhead and avoid OOMs in reading data.")
.intConf
.createWithDefault(4096)

val ORC_COPY_BATCH_TO_SPARK = buildConf("spark.sql.orc.copyBatchToSpark")
.doc("Whether or not to copy the ORC columnar batch to Spark columnar batch in the " +
"vectorized ORC reader.")
Expand Down Expand Up @@ -1216,10 +1228,14 @@ class SQLConf extends Serializable with Logging {

def orcVectorizedReaderEnabled: Boolean = getConf(ORC_VECTORIZED_READER_ENABLED)

def orcVectorizedReaderBatchSize: Int = getConf(ORC_VECTORIZED_READER_BATCH_SIZE)

def parquetCompressionCodec: String = getConf(PARQUET_COMPRESSION)

def parquetVectorizedReaderEnabled: Boolean = getConf(PARQUET_VECTORIZED_READER_ENABLED)

def parquetVectorizedReaderBatchSize: Int = getConf(PARQUET_VECTORIZED_READER_BATCH_SIZE)

def columnBatchSize: Int = getConf(COLUMN_BATCH_SIZE)

def numShufflePartitions: Int = getConf(SHUFFLE_PARTITIONS)
Expand Down
22 changes: 12 additions & 10 deletions .../src/main/java/org/apache/spark/sql/execution/datasources/orc/OrcColumnarBatchReader.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,8 +49,9 @@
* After creating, `initialize` and `initBatch` should be called sequentially.
*/
public class OrcColumnarBatchReader extends RecordReader<Void, ColumnarBatch> {
// TODO: make this configurable.
private static final int CAPACITY = 4 * 1024;

// The default size of vectorized batch.

@cloud-fan cloud-fan Jan 30, 2018

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe we can remove the comment. It's just the capacity, not a default value.

@jiangxb1987 jiangxb1987 Feb 1, 2018

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about rephrase to The capacity of vectorized batch ?

private int capacity;

// Vectorized ORC Row Batch
private VectorizedRowBatch batch;
Expand Down Expand Up @@ -81,9 +82,10 @@ public class OrcColumnarBatchReader extends RecordReader<Void, ColumnarBatch> {
// Whether or not to copy the ORC columnar batch to Spark columnar batch.
private final boolean copyToSpark;

public OrcColumnarBatchReader(boolean useOffHeap, boolean copyToSpark) {
public OrcColumnarBatchReader(boolean useOffHeap, boolean copyToSpark, int capacity) {
MEMORY_MODE = useOffHeap ? MemoryMode.OFF_HEAP : MemoryMode.ON_HEAP;
this.copyToSpark = copyToSpark;
this.capacity = capacity;
}


Expand Down Expand Up @@ -148,7 +150,7 @@ public void initBatch(
StructField[] requiredFields,
StructType partitionSchema,
InternalRow partitionValues) {
batch = orcSchema.createRowBatch(CAPACITY);
batch = orcSchema.createRowBatch(capacity);
assert(!batch.selectedInUse); // `selectedInUse` should be initialized with `false`.

this.requiredFields = requiredFields;
Expand All @@ -162,15 +164,15 @@ public void initBatch(

if (copyToSpark) {
if (MEMORY_MODE == MemoryMode.OFF_HEAP) {
columnVectors = OffHeapColumnVector.allocateColumns(CAPACITY, resultSchema);
columnVectors = OffHeapColumnVector.allocateColumns(capacity, resultSchema);
} else {
columnVectors = OnHeapColumnVector.allocateColumns(CAPACITY, resultSchema);
columnVectors = OnHeapColumnVector.allocateColumns(capacity, resultSchema);
}

// Initialize the missing columns once.
for (int i = 0; i < requiredFields.length; i++) {
if (requestedColIds[i] == -1) {
columnVectors[i].putNulls(0, CAPACITY);
columnVectors[i].putNulls(0, capacity);
columnVectors[i].setIsConstant();
}
}
Expand All @@ -193,8 +195,8 @@ public void initBatch(
int colId = requestedColIds[i];
// Initialize the missing columns once.
if (colId == -1) {
OnHeapColumnVector missingCol = new OnHeapColumnVector(CAPACITY, dt);
missingCol.putNulls(0, CAPACITY);
OnHeapColumnVector missingCol = new OnHeapColumnVector(capacity, dt);
missingCol.putNulls(0, capacity);
missingCol.setIsConstant();
orcVectorWrappers[i] = missingCol;
} else {
Expand All @@ -206,7 +208,7 @@ public void initBatch(
int partitionIdx = requiredFields.length;
for (int i = 0; i < partitionValues.numFields(); i++) {
DataType dt = partitionSchema.fields()[i].dataType();
OnHeapColumnVector partitionCol = new OnHeapColumnVector(CAPACITY, dt);
OnHeapColumnVector partitionCol = new OnHeapColumnVector(capacity, dt);
ColumnVectorUtils.populate(partitionCol, partitionValues, i);
partitionCol.setIsConstant();
orcVectorWrappers[partitionIdx + i] = partitionCol;
Expand Down
19 changes: 11 additions & 8 deletions ...ava/org/apache/spark/sql/execution/datasources/parquet/VectorizedParquetRecordReader.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,8 +50,9 @@
* TODO: make this always return ColumnarBatches.
*/
public class VectorizedParquetRecordReader extends SpecificParquetRecordReaderBase<Object> {
// TODO: make this configurable.
private static final int CAPACITY = 4 * 1024;

// The default size of vectorized batch.
private int capacity;

/**
* Batch of rows that we assemble and the current index we've returned. Every time this
Expand Down Expand Up @@ -115,13 +116,15 @@ public class VectorizedParquetRecordReader extends SpecificParquetRecordReaderBa
*/
private final MemoryMode MEMORY_MODE;

public VectorizedParquetRecordReader(TimeZone convertTz, boolean useOffHeap) {
public VectorizedParquetRecordReader(TimeZone convertTz, boolean useOffHeap, int capacity) {
this.convertTz = convertTz;
MEMORY_MODE = useOffHeap ? MemoryMode.OFF_HEAP : MemoryMode.ON_HEAP;
this.capacity = capacity;
}

// Vectorized parquet reader used for testing and benchmark.
public VectorizedParquetRecordReader(boolean useOffHeap) {
this(null, useOffHeap);
this(null, useOffHeap, 4096);

@kiszk kiszk Jan 25, 2018

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about changing benchmark and test programs to pass capacity and remove this constructor?
These programs also have SQLConf.

@cloud-fan cloud-fan Jan 30, 2018

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's good to avoid hardcoding the default value again in the code. If there are only a few places need to be changed, let's do it.

}

/**
Expand Down Expand Up @@ -199,9 +202,9 @@ private void initBatch(
}

if (memMode == MemoryMode.OFF_HEAP) {
columnVectors = OffHeapColumnVector.allocateColumns(CAPACITY, batchSchema);
columnVectors = OffHeapColumnVector.allocateColumns(capacity, batchSchema);
} else {
columnVectors = OnHeapColumnVector.allocateColumns(CAPACITY, batchSchema);
columnVectors = OnHeapColumnVector.allocateColumns(capacity, batchSchema);
}
columnarBatch = new ColumnarBatch(columnVectors);
if (partitionColumns != null) {
Expand All @@ -215,7 +218,7 @@ private void initBatch(
// Initialize missing columns with nulls.
for (int i = 0; i < missingColumns.length; i++) {
if (missingColumns[i]) {
columnVectors[i].putNulls(0, CAPACITY);
columnVectors[i].putNulls(0, capacity);
columnVectors[i].setIsConstant();
}
}
Expand Down Expand Up @@ -257,7 +260,7 @@ public boolean nextBatch() throws IOException {
if (rowsReturned >= totalRowCount) return false;
checkEndOfRowGroup();

int num = (int) Math.min((long) CAPACITY, totalCountLoadedSoFar - rowsReturned);
int num = (int) Math.min((long) capacity, totalCountLoadedSoFar - rowsReturned);
for (int i = 0; i < columnReaders.length; ++i) {
if (columnReaders[i] == null) continue;
columnReaders[i].readBatch(num, columnVectors[i]);
Expand Down
7 changes: 5 additions & 2 deletions sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/WritableColumnVector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,8 +96,11 @@ public void reserve(int requiredCapacity) {
private void throwUnsupportedException(int requiredCapacity, Throwable cause) {
String message = "Cannot reserve additional contiguous bytes in the vectorized reader " +
"(requested = " + requiredCapacity + " bytes). As a workaround, you can disable the " +
"vectorized reader by setting " + SQLConf.PARQUET_VECTORIZED_READER_ENABLED().key() +
" to false.";
"vectorized reader, or increase the vectorized reader batch size. For parquet file " +
"format, refer to " + SQLConf.PARQUET_VECTORIZED_READER_ENABLED().key() + " and " +
SQLConf.PARQUET_VECTORIZED_READER_BATCH_SIZE().key() + "; for orc file format, refer to " +
SQLConf.ORC_VECTORIZED_READER_ENABLED().key() + " and " +
SQLConf.ORC_VECTORIZED_READER_BATCH_SIZE().key() + ".";
throw new RuntimeException(message, cause);
}

Expand Down
3 changes: 2 additions & 1 deletion sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcFileFormat.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,7 @@ class OrcFileFormat
val sqlConf = sparkSession.sessionState.conf
val enableOffHeapColumnVector = sqlConf.offHeapColumnVectorEnabled
val enableVectorizedReader = supportBatch(sparkSession, resultSchema)
val capacity = sqlConf.orcVectorizedReaderBatchSize
val copyToSpark = sparkSession.sessionState.conf.getConf(SQLConf.ORC_COPY_BATCH_TO_SPARK)

val broadcastedConf =
Expand Down Expand Up @@ -186,7 +187,7 @@ class OrcFileFormat
val taskContext = Option(TaskContext.get())
if (enableVectorizedReader) {
val batchReader = new OrcColumnarBatchReader(
enableOffHeapColumnVector && taskContext.isDefined, copyToSpark)
enableOffHeapColumnVector && taskContext.isDefined, copyToSpark, capacity)
batchReader.initialize(fileSplit, taskAttemptContext)
batchReader.initBatch(
reader.getSchema,
Expand Down
3 changes: 2 additions & 1 deletion ...src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFileFormat.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -350,6 +350,7 @@ class ParquetFileFormat
sparkSession.sessionState.conf.parquetRecordFilterEnabled
val timestampConversion: Boolean =
sparkSession.sessionState.conf.isParquetINT96TimestampConversion
val capacity = sqlConf.parquetVectorizedReaderBatchSize
// Whole stage codegen (PhysicalRDD) is able to deal with batches directly
val returningBatch = supportBatch(sparkSession, resultSchema)

Expand Down Expand Up @@ -396,7 +397,7 @@ class ParquetFileFormat
val taskContext = Option(TaskContext.get())
val parquetReader = if (enableVectorizedReader) {
val vectorizedReader = new VectorizedParquetRecordReader(
convertTz.orNull, enableOffHeapColumnVector && taskContext.isDefined)
convertTz.orNull, enableOffHeapColumnVector && taskContext.isDefined, capacity)
vectorizedReader.initialize(split, hadoopAttemptContext)
logDebug(s"Appending $partitionSchema ${file.partitionValues}")
vectorizedReader.initBatch(partitionSchema, file.partitionValues)
Expand Down