fix flaky tests

apache · Oct 7, 2016 · 52b2b7d · 52b2b7d
1 parent 566d7f2
commit 52b2b7d
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ColumnTypeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ColumnTypeSuite.scala
@@ -101,14 +101,15 @@ class ColumnTypeSuite extends SparkFunSuite with Logging {
 
   def testColumnType[JvmType](columnType: ColumnType[JvmType]): Unit = {
 
-    val buffer = ByteBuffer.allocate(DEFAULT_BUFFER_SIZE).order(ByteOrder.nativeOrder())
     val proj = UnsafeProjection.create(Array[DataType](columnType.dataType))
     val converter = CatalystTypeConverters.createToScalaConverter(columnType.dataType)
     val seq = (0 until 4).map(_ => proj(makeRandomRow(columnType)).copy())
+    val totalSize = seq.map(_.getSizeInBytes).sum
+    val bufferSize = Math.max(DEFAULT_BUFFER_SIZE, totalSize)
 
     test(s"$columnType append/extract") {
-      buffer.rewind()
-      seq.foreach(columnType.append(_, 0, buffer))
+      val buffer = ByteBuffer.allocate(bufferSize).order(ByteOrder.nativeOrder())
+      seq.foreach(r => columnType.append(columnType.getField(r, 0), buffer))
 
       buffer.rewind()
       seq.foreach { row =>