You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by li...@apache.org on 2018/01/19 16:58:25 UTC
spark git commit: [SPARK-23149][SQL] polish ColumnarBatch
Repository: spark
Updated Branches:
refs/heads/master 606a7485f -> d8aaa771e
[SPARK-23149][SQL] polish ColumnarBatch
## What changes were proposed in this pull request?
Several cleanups in `ColumnarBatch`
* remove `schema`. The `ColumnVector`s inside `ColumnarBatch` already have the data type information, we don't need this `schema`.
* remove `capacity`. `ColumnarBatch` is just a wrapper of `ColumnVector`s, not builders, it doesn't need a capacity property.
* remove `DEFAULT_BATCH_SIZE`. As a wrapper, `ColumnarBatch` can't decide the batch size, it should be decided by the reader, e.g. parquet reader, orc reader, cached table reader. The default batch size should also be defined by the reader.
## How was this patch tested?
existing tests.
Author: Wenchen Fan <we...@databricks.com>
Closes #20316 from cloud-fan/columnar-batch.
Project: http://git-wip-us.apache.org/repos/asf/spark/repo
Commit: http://git-wip-us.apache.org/repos/asf/spark/commit/d8aaa771
Tree: http://git-wip-us.apache.org/repos/asf/spark/tree/d8aaa771
Diff: http://git-wip-us.apache.org/repos/asf/spark/diff/d8aaa771
Branch: refs/heads/master
Commit: d8aaa771e249b3f54b57ce24763e53fd65a0dbf7
Parents: 606a748
Author: Wenchen Fan <we...@databricks.com>
Authored: Fri Jan 19 08:58:21 2018 -0800
Committer: gatorsmile <ga...@gmail.com>
Committed: Fri Jan 19 08:58:21 2018 -0800
----------------------------------------------------------------------
.../datasources/orc/OrcColumnarBatchReader.java | 49 +++++++-------------
.../SpecificParquetRecordReaderBase.java | 12 ++---
.../parquet/VectorizedParquetRecordReader.java | 24 +++++-----
.../execution/vectorized/ColumnVectorUtils.java | 18 +++----
.../spark/sql/vectorized/ColumnarBatch.java | 20 +-------
.../aggregate/VectorizedHashMapGenerator.scala | 2 +-
.../sql/execution/arrow/ArrowConverters.scala | 2 +-
.../columnar/InMemoryTableScanExec.scala | 5 +-
.../execution/python/ArrowEvalPythonExec.scala | 8 ++--
.../execution/python/ArrowPythonRunner.scala | 2 +-
.../sql/sources/v2/JavaBatchDataSourceV2.java | 3 +-
.../vectorized/ColumnarBatchSuite.scala | 7 ++-
.../sql/sources/v2/DataSourceV2Suite.scala | 3 +-
13 files changed, 61 insertions(+), 94 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/spark/blob/d8aaa771/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/orc/OrcColumnarBatchReader.java
----------------------------------------------------------------------
diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/orc/OrcColumnarBatchReader.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/orc/OrcColumnarBatchReader.java
index 36fdf2b..89bae43 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/orc/OrcColumnarBatchReader.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/orc/OrcColumnarBatchReader.java
@@ -49,18 +49,8 @@ import org.apache.spark.sql.vectorized.ColumnarBatch;
* After creating, `initialize` and `initBatch` should be called sequentially.
*/
public class OrcColumnarBatchReader extends RecordReader<Void, ColumnarBatch> {
-
- /**
- * The default size of batch. We use this value for ORC reader to make it consistent with Spark's
- * columnar batch, because their default batch sizes are different like the following:
- *
- * - ORC's VectorizedRowBatch.DEFAULT_SIZE = 1024
- * - Spark's ColumnarBatch.DEFAULT_BATCH_SIZE = 4 * 1024
- */
- private static final int DEFAULT_SIZE = 4 * 1024;
-
- // ORC File Reader
- private Reader reader;
+ // TODO: make this configurable.
+ private static final int CAPACITY = 4 * 1024;
// Vectorized ORC Row Batch
private VectorizedRowBatch batch;
@@ -98,22 +88,22 @@ public class OrcColumnarBatchReader extends RecordReader<Void, ColumnarBatch> {
@Override
- public Void getCurrentKey() throws IOException, InterruptedException {
+ public Void getCurrentKey() {
return null;
}
@Override
- public ColumnarBatch getCurrentValue() throws IOException, InterruptedException {
+ public ColumnarBatch getCurrentValue() {
return columnarBatch;
}
@Override
- public float getProgress() throws IOException, InterruptedException {
+ public float getProgress() throws IOException {
return recordReader.getProgress();
}
@Override
- public boolean nextKeyValue() throws IOException, InterruptedException {
+ public boolean nextKeyValue() throws IOException {
return nextBatch();
}
@@ -134,16 +124,15 @@ public class OrcColumnarBatchReader extends RecordReader<Void, ColumnarBatch> {
* Please note that `initBatch` is needed to be called after this.
*/
@Override
- public void initialize(InputSplit inputSplit, TaskAttemptContext taskAttemptContext)
- throws IOException, InterruptedException {
+ public void initialize(
+ InputSplit inputSplit, TaskAttemptContext taskAttemptContext) throws IOException {
FileSplit fileSplit = (FileSplit)inputSplit;
Configuration conf = taskAttemptContext.getConfiguration();
- reader = OrcFile.createReader(
+ Reader reader = OrcFile.createReader(
fileSplit.getPath(),
OrcFile.readerOptions(conf)
.maxLength(OrcConf.MAX_FILE_LENGTH.getLong(conf))
.filesystem(fileSplit.getPath().getFileSystem(conf)));
-
Reader.Options options =
OrcInputFormat.buildOptions(conf, reader, fileSplit.getStart(), fileSplit.getLength());
recordReader = reader.rows(options);
@@ -159,7 +148,7 @@ public class OrcColumnarBatchReader extends RecordReader<Void, ColumnarBatch> {
StructField[] requiredFields,
StructType partitionSchema,
InternalRow partitionValues) {
- batch = orcSchema.createRowBatch(DEFAULT_SIZE);
+ batch = orcSchema.createRowBatch(CAPACITY);
assert(!batch.selectedInUse); // `selectedInUse` should be initialized with `false`.
this.requiredFields = requiredFields;
@@ -171,19 +160,17 @@ public class OrcColumnarBatchReader extends RecordReader<Void, ColumnarBatch> {
resultSchema = resultSchema.add(f);
}
- int capacity = DEFAULT_SIZE;
-
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();
}
}
@@ -196,7 +183,7 @@ public class OrcColumnarBatchReader extends RecordReader<Void, ColumnarBatch> {
}
}
- columnarBatch = new ColumnarBatch(resultSchema, columnVectors, capacity);
+ columnarBatch = new ColumnarBatch(columnVectors);
} else {
// Just wrap the ORC column vector instead of copying it to Spark column vector.
orcVectorWrappers = new org.apache.spark.sql.vectorized.ColumnVector[resultSchema.length()];
@@ -206,8 +193,8 @@ public class OrcColumnarBatchReader extends RecordReader<Void, ColumnarBatch> {
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 {
@@ -219,14 +206,14 @@ public class OrcColumnarBatchReader extends RecordReader<Void, ColumnarBatch> {
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;
}
}
- columnarBatch = new ColumnarBatch(resultSchema, orcVectorWrappers, capacity);
+ columnarBatch = new ColumnarBatch(orcVectorWrappers);
}
}
http://git-wip-us.apache.org/repos/asf/spark/blob/d8aaa771/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/SpecificParquetRecordReaderBase.java
----------------------------------------------------------------------
diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/SpecificParquetRecordReaderBase.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/SpecificParquetRecordReaderBase.java
index 80c2f49..e65cd25 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/SpecificParquetRecordReaderBase.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/SpecificParquetRecordReaderBase.java
@@ -170,7 +170,7 @@ public abstract class SpecificParquetRecordReaderBase<T> extends RecordReader<Vo
* Returns the list of files at 'path' recursively. This skips files that are ignored normally
* by MapReduce.
*/
- public static List<String> listDirectory(File path) throws IOException {
+ public static List<String> listDirectory(File path) {
List<String> result = new ArrayList<>();
if (path.isDirectory()) {
for (File f: path.listFiles()) {
@@ -231,7 +231,7 @@ public abstract class SpecificParquetRecordReaderBase<T> extends RecordReader<Vo
}
@Override
- public Void getCurrentKey() throws IOException, InterruptedException {
+ public Void getCurrentKey() {
return null;
}
@@ -259,7 +259,7 @@ public abstract class SpecificParquetRecordReaderBase<T> extends RecordReader<Vo
}
@Override
- int nextInt() throws IOException {
+ int nextInt() {
return delegate.readInteger();
}
}
@@ -279,15 +279,15 @@ public abstract class SpecificParquetRecordReaderBase<T> extends RecordReader<Vo
protected static final class NullIntIterator extends IntIterator {
@Override
- int nextInt() throws IOException { return 0; }
+ int nextInt() { return 0; }
}
/**
* Creates a reader for definition and repetition levels, returning an optimized one if
* the levels are not needed.
*/
- protected static IntIterator createRLEIterator(int maxLevel, BytesInput bytes,
- ColumnDescriptor descriptor) throws IOException {
+ protected static IntIterator createRLEIterator(
+ int maxLevel, BytesInput bytes, ColumnDescriptor descriptor) throws IOException {
try {
if (maxLevel == 0) return new NullIntIterator();
return new RLEIntIterator(
http://git-wip-us.apache.org/repos/asf/spark/blob/d8aaa771/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedParquetRecordReader.java
----------------------------------------------------------------------
diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedParquetRecordReader.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedParquetRecordReader.java
index cd745b1..bb1b236 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedParquetRecordReader.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedParquetRecordReader.java
@@ -50,6 +50,9 @@ import org.apache.spark.sql.types.StructType;
* TODO: make this always return ColumnarBatches.
*/
public class VectorizedParquetRecordReader extends SpecificParquetRecordReaderBase<Object> {
+ // TODO: make this configurable.
+ private static final int CAPACITY = 4 * 1024;
+
/**
* Batch of rows that we assemble and the current index we've returned. Every time this
* batch is used up (batchIdx == numBatched), we populated the batch.
@@ -152,7 +155,7 @@ public class VectorizedParquetRecordReader extends SpecificParquetRecordReaderBa
}
@Override
- public boolean nextKeyValue() throws IOException, InterruptedException {
+ public boolean nextKeyValue() throws IOException {
resultBatch();
if (returnColumnarBatch) return nextBatch();
@@ -165,13 +168,13 @@ public class VectorizedParquetRecordReader extends SpecificParquetRecordReaderBa
}
@Override
- public Object getCurrentValue() throws IOException, InterruptedException {
+ public Object getCurrentValue() {
if (returnColumnarBatch) return columnarBatch;
return columnarBatch.getRow(batchIdx - 1);
}
@Override
- public float getProgress() throws IOException, InterruptedException {
+ public float getProgress() {
return (float) rowsReturned / totalRowCount;
}
@@ -181,7 +184,7 @@ public class VectorizedParquetRecordReader extends SpecificParquetRecordReaderBa
// Columns 0,1: data columns
// Column 2: partitionValues[0]
// Column 3: partitionValues[1]
- public void initBatch(
+ private void initBatch(
MemoryMode memMode,
StructType partitionColumns,
InternalRow partitionValues) {
@@ -195,13 +198,12 @@ public class VectorizedParquetRecordReader extends SpecificParquetRecordReaderBa
}
}
- int capacity = ColumnarBatch.DEFAULT_BATCH_SIZE;
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(batchSchema, columnVectors, capacity);
+ columnarBatch = new ColumnarBatch(columnVectors);
if (partitionColumns != null) {
int partitionIdx = sparkSchema.fields().length;
for (int i = 0; i < partitionColumns.fields().length; i++) {
@@ -213,13 +215,13 @@ public class VectorizedParquetRecordReader extends SpecificParquetRecordReaderBa
// Initialize missing columns with nulls.
for (int i = 0; i < missingColumns.length; i++) {
if (missingColumns[i]) {
- columnVectors[i].putNulls(0, columnarBatch.capacity());
+ columnVectors[i].putNulls(0, CAPACITY);
columnVectors[i].setIsConstant();
}
}
}
- public void initBatch() {
+ private void initBatch() {
initBatch(MEMORY_MODE, null, null);
}
@@ -255,7 +257,7 @@ public class VectorizedParquetRecordReader extends SpecificParquetRecordReaderBa
if (rowsReturned >= totalRowCount) return false;
checkEndOfRowGroup();
- int num = (int) Math.min((long) columnarBatch.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]);
http://git-wip-us.apache.org/repos/asf/spark/blob/d8aaa771/sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ColumnVectorUtils.java
----------------------------------------------------------------------
diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ColumnVectorUtils.java b/sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ColumnVectorUtils.java
index b5cbe8e..5ee8cc8 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ColumnVectorUtils.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ColumnVectorUtils.java
@@ -118,19 +118,19 @@ public class ColumnVectorUtils {
}
} else {
if (t == DataTypes.BooleanType) {
- dst.appendBoolean(((Boolean)o).booleanValue());
+ dst.appendBoolean((Boolean) o);
} else if (t == DataTypes.ByteType) {
- dst.appendByte(((Byte) o).byteValue());
+ dst.appendByte((Byte) o);
} else if (t == DataTypes.ShortType) {
- dst.appendShort(((Short)o).shortValue());
+ dst.appendShort((Short) o);
} else if (t == DataTypes.IntegerType) {
- dst.appendInt(((Integer)o).intValue());
+ dst.appendInt((Integer) o);
} else if (t == DataTypes.LongType) {
- dst.appendLong(((Long)o).longValue());
+ dst.appendLong((Long) o);
} else if (t == DataTypes.FloatType) {
- dst.appendFloat(((Float)o).floatValue());
+ dst.appendFloat((Float) o);
} else if (t == DataTypes.DoubleType) {
- dst.appendDouble(((Double)o).doubleValue());
+ dst.appendDouble((Double) o);
} else if (t == DataTypes.StringType) {
byte[] b =((String)o).getBytes(StandardCharsets.UTF_8);
dst.appendByteArray(b, 0, b.length);
@@ -192,7 +192,7 @@ public class ColumnVectorUtils {
*/
public static ColumnarBatch toBatch(
StructType schema, MemoryMode memMode, Iterator<Row> row) {
- int capacity = ColumnarBatch.DEFAULT_BATCH_SIZE;
+ int capacity = 4 * 1024;
WritableColumnVector[] columnVectors;
if (memMode == MemoryMode.OFF_HEAP) {
columnVectors = OffHeapColumnVector.allocateColumns(capacity, schema);
@@ -208,7 +208,7 @@ public class ColumnVectorUtils {
}
n++;
}
- ColumnarBatch batch = new ColumnarBatch(schema, columnVectors, capacity);
+ ColumnarBatch batch = new ColumnarBatch(columnVectors);
batch.setNumRows(n);
return batch;
}
http://git-wip-us.apache.org/repos/asf/spark/blob/d8aaa771/sql/core/src/main/java/org/apache/spark/sql/vectorized/ColumnarBatch.java
----------------------------------------------------------------------
diff --git a/sql/core/src/main/java/org/apache/spark/sql/vectorized/ColumnarBatch.java b/sql/core/src/main/java/org/apache/spark/sql/vectorized/ColumnarBatch.java
index 9ae1c6d..4dc826c 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/vectorized/ColumnarBatch.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/vectorized/ColumnarBatch.java
@@ -20,7 +20,6 @@ import java.util.*;
import org.apache.spark.sql.catalyst.InternalRow;
import org.apache.spark.sql.execution.vectorized.MutableColumnarRow;
-import org.apache.spark.sql.types.StructType;
/**
* This class wraps multiple ColumnVectors as a row-wise table. It provides a row view of this
@@ -28,10 +27,6 @@ import org.apache.spark.sql.types.StructType;
* the entire data loading process.
*/
public final class ColumnarBatch {
- public static final int DEFAULT_BATCH_SIZE = 4 * 1024;
-
- private final StructType schema;
- private final int capacity;
private int numRows;
private final ColumnVector[] columns;
@@ -82,7 +77,6 @@ public final class ColumnarBatch {
* Sets the number of rows in this batch.
*/
public void setNumRows(int numRows) {
- assert(numRows <= this.capacity);
this.numRows = numRows;
}
@@ -97,16 +91,6 @@ public final class ColumnarBatch {
public int numRows() { return numRows; }
/**
- * Returns the schema that makes up this batch.
- */
- public StructType schema() { return schema; }
-
- /**
- * Returns the max capacity (in number of rows) for this batch.
- */
- public int capacity() { return capacity; }
-
- /**
* Returns the column at `ordinal`.
*/
public ColumnVector column(int ordinal) { return columns[ordinal]; }
@@ -120,10 +104,8 @@ public final class ColumnarBatch {
return row;
}
- public ColumnarBatch(StructType schema, ColumnVector[] columns, int capacity) {
- this.schema = schema;
+ public ColumnarBatch(ColumnVector[] columns) {
this.columns = columns;
- this.capacity = capacity;
this.row = new MutableColumnarRow(columns);
}
}
http://git-wip-us.apache.org/repos/asf/spark/blob/d8aaa771/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/VectorizedHashMapGenerator.scala
----------------------------------------------------------------------
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/VectorizedHashMapGenerator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/VectorizedHashMapGenerator.scala
index 0cf9b53..eb48584 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/VectorizedHashMapGenerator.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/VectorizedHashMapGenerator.scala
@@ -94,7 +94,7 @@ class VectorizedHashMapGenerator(
|
| public $generatedClassName() {
| vectors = ${classOf[OnHeapColumnVector].getName}.allocateColumns(capacity, schema);
- | batch = new ${classOf[ColumnarBatch].getName}(schema, vectors, capacity);
+ | batch = new ${classOf[ColumnarBatch].getName}(vectors);
|
| // Generates a projection to return the aggregate buffer only.
| ${classOf[OnHeapColumnVector].getName}[] aggBufferVectors =
http://git-wip-us.apache.org/repos/asf/spark/blob/d8aaa771/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala
----------------------------------------------------------------------
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala
index bcd1aa0..7487564 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala
@@ -175,7 +175,7 @@ private[sql] object ArrowConverters {
new ArrowColumnVector(vector).asInstanceOf[ColumnVector]
}.toArray
- val batch = new ColumnarBatch(schemaRead, columns, root.getRowCount)
+ val batch = new ColumnarBatch(columns)
batch.setNumRows(root.getRowCount)
batch.rowIterator().asScala
}
http://git-wip-us.apache.org/repos/asf/spark/blob/d8aaa771/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryTableScanExec.scala
----------------------------------------------------------------------
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryTableScanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryTableScanExec.scala
index 3565ee3..28b3875 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryTableScanExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryTableScanExec.scala
@@ -78,11 +78,10 @@ case class InMemoryTableScanExec(
} else {
OffHeapColumnVector.allocateColumns(rowCount, columnarBatchSchema)
}
- val columnarBatch = new ColumnarBatch(
- columnarBatchSchema, columnVectors.asInstanceOf[Array[ColumnVector]], rowCount)
+ val columnarBatch = new ColumnarBatch(columnVectors.asInstanceOf[Array[ColumnVector]])
columnarBatch.setNumRows(rowCount)
- for (i <- 0 until attributes.length) {
+ for (i <- attributes.indices) {
ColumnAccessor.decompress(
cachedColumnarBatch.buffers(columnIndices(i)),
columnarBatch.column(i).asInstanceOf[WritableColumnVector],
http://git-wip-us.apache.org/repos/asf/spark/blob/d8aaa771/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala
----------------------------------------------------------------------
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala
index c06bc7b..47b146f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala
@@ -74,8 +74,7 @@ case class ArrowEvalPythonExec(udfs: Seq[PythonUDF], output: Seq[Attribute], chi
schema: StructType,
context: TaskContext): Iterator[InternalRow] = {
- val schemaOut = StructType.fromAttributes(output.drop(child.output.length).zipWithIndex
- .map { case (attr, i) => attr.withName(s"_$i") })
+ val outputTypes = output.drop(child.output.length).map(_.dataType)
// DO NOT use iter.grouped(). See BatchIterator.
val batchIter = if (batchSize > 0) new BatchIterator(iter, batchSize) else Iterator(iter)
@@ -90,8 +89,9 @@ case class ArrowEvalPythonExec(udfs: Seq[PythonUDF], output: Seq[Attribute], chi
private var currentIter = if (columnarBatchIter.hasNext) {
val batch = columnarBatchIter.next()
- assert(schemaOut.equals(batch.schema),
- s"Invalid schema from pandas_udf: expected $schemaOut, got ${batch.schema}")
+ val actualDataTypes = (0 until batch.numCols()).map(i => batch.column(i).dataType())
+ assert(outputTypes == actualDataTypes, "Invalid schema from pandas_udf: " +
+ s"expected ${outputTypes.mkString(", ")}, got ${actualDataTypes.mkString(", ")}")
batch.rowIterator.asScala
} else {
Iterator.empty
http://git-wip-us.apache.org/repos/asf/spark/blob/d8aaa771/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
----------------------------------------------------------------------
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
index dc5ba96..5fcdcdd 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
@@ -138,7 +138,7 @@ class ArrowPythonRunner(
if (reader != null && batchLoaded) {
batchLoaded = reader.loadNextBatch()
if (batchLoaded) {
- val batch = new ColumnarBatch(schema, vectors, root.getRowCount)
+ val batch = new ColumnarBatch(vectors)
batch.setNumRows(root.getRowCount)
batch
} else {
http://git-wip-us.apache.org/repos/asf/spark/blob/d8aaa771/sql/core/src/test/java/test/org/apache/spark/sql/sources/v2/JavaBatchDataSourceV2.java
----------------------------------------------------------------------
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/sources/v2/JavaBatchDataSourceV2.java b/sql/core/src/test/java/test/org/apache/spark/sql/sources/v2/JavaBatchDataSourceV2.java
index 98d6a53..a5d77a9 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/sources/v2/JavaBatchDataSourceV2.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/sources/v2/JavaBatchDataSourceV2.java
@@ -69,8 +69,7 @@ public class JavaBatchDataSourceV2 implements DataSourceV2, ReadSupport {
ColumnVector[] vectors = new ColumnVector[2];
vectors[0] = i;
vectors[1] = j;
- this.batch =
- new ColumnarBatch(new StructType().add("i", "int").add("j", "int"), vectors, BATCH_SIZE);
+ this.batch = new ColumnarBatch(vectors);
return this;
}
http://git-wip-us.apache.org/repos/asf/spark/blob/d8aaa771/sql/core/src/test/scala/org/apache/spark/sql/execution/vectorized/ColumnarBatchSuite.scala
----------------------------------------------------------------------
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/vectorized/ColumnarBatchSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/vectorized/ColumnarBatchSuite.scala
index 675f06b..cd90681 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/vectorized/ColumnarBatchSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/vectorized/ColumnarBatchSuite.scala
@@ -875,14 +875,13 @@ class ColumnarBatchSuite extends SparkFunSuite {
.add("intCol2", IntegerType)
.add("string", BinaryType)
- val capacity = ColumnarBatch.DEFAULT_BATCH_SIZE
+ val capacity = 4 * 1024
val columns = schema.fields.map { field =>
allocate(capacity, field.dataType, memMode)
}
- val batch = new ColumnarBatch(schema, columns.toArray, ColumnarBatch.DEFAULT_BATCH_SIZE)
+ val batch = new ColumnarBatch(columns.toArray)
assert(batch.numCols() == 4)
assert(batch.numRows() == 0)
- assert(batch.capacity() > 0)
assert(batch.rowIterator().hasNext == false)
// Add a row [1, 1.1, NULL]
@@ -1153,7 +1152,7 @@ class ColumnarBatchSuite extends SparkFunSuite {
val columnVectors = Seq(new ArrowColumnVector(vector1), new ArrowColumnVector(vector2))
val schema = StructType(Seq(StructField("int1", IntegerType), StructField("int2", IntegerType)))
- val batch = new ColumnarBatch(schema, columnVectors.toArray[ColumnVector], 11)
+ val batch = new ColumnarBatch(columnVectors.toArray)
batch.setNumRows(11)
assert(batch.numCols() == 2)
http://git-wip-us.apache.org/repos/asf/spark/blob/d8aaa771/sql/core/src/test/scala/org/apache/spark/sql/sources/v2/DataSourceV2Suite.scala
----------------------------------------------------------------------
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/v2/DataSourceV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/v2/DataSourceV2Suite.scala
index a89f7c5..0ca2952 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/sources/v2/DataSourceV2Suite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/v2/DataSourceV2Suite.scala
@@ -333,8 +333,7 @@ class BatchReadTask(start: Int, end: Int)
private final val BATCH_SIZE = 20
private lazy val i = new OnHeapColumnVector(BATCH_SIZE, IntegerType)
private lazy val j = new OnHeapColumnVector(BATCH_SIZE, IntegerType)
- private lazy val batch = new ColumnarBatch(
- new StructType().add("i", "int").add("j", "int"), Array(i, j), BATCH_SIZE)
+ private lazy val batch = new ColumnarBatch(Array(i, j))
private var current = start
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org