@@ -50,19 +50,19 @@ class ColumnarBatchSuite extends SparkFunSuite {
50
50
name : String ,
51
51
size : Int ,
52
52
dt : DataType )(
53
- block : ( WritableColumnVector , MemoryMode ) => Unit ): Unit = {
53
+ block : WritableColumnVector => Unit ): Unit = {
54
54
test(name) {
55
55
Seq (MemoryMode .ON_HEAP , MemoryMode .OFF_HEAP ).foreach { mode =>
56
56
val vector = allocate(size, dt, mode)
57
- try block(vector, mode ) finally {
57
+ try block(vector) finally {
58
58
vector.close()
59
59
}
60
60
}
61
61
}
62
62
}
63
63
64
64
testVector(" Null APIs" , 1024 , IntegerType ) {
65
- ( column, memMode) =>
65
+ column =>
66
66
val reference = mutable.ArrayBuffer .empty[Boolean ]
67
67
var idx = 0
68
68
assert(! column.anyNullsSet())
@@ -121,15 +121,11 @@ class ColumnarBatchSuite extends SparkFunSuite {
121
121
122
122
reference.zipWithIndex.foreach { v =>
123
123
assert(v._1 == column.isNullAt(v._2))
124
- if (memMode == MemoryMode .OFF_HEAP ) {
125
- val addr = column.nullsNativeAddress()
126
- assert(v._1 == (Platform .getByte(null , addr + v._2) == 1 ), " index=" + v._2)
127
- }
128
124
}
129
125
}
130
126
131
127
testVector(" Byte APIs" , 1024 , ByteType ) {
132
- ( column, memMode) =>
128
+ column =>
133
129
val reference = mutable.ArrayBuffer .empty[Byte ]
134
130
135
131
var values = (10 :: 20 :: 30 :: 40 :: 50 :: Nil ).map(_.toByte).toArray
@@ -173,16 +169,12 @@ class ColumnarBatchSuite extends SparkFunSuite {
173
169
idx += 3
174
170
175
171
reference.zipWithIndex.foreach { v =>
176
- assert(v._1 == column.getByte(v._2), " MemoryMode" + memMode)
177
- if (memMode == MemoryMode .OFF_HEAP ) {
178
- val addr = column.valuesNativeAddress()
179
- assert(v._1 == Platform .getByte(null , addr + v._2))
180
- }
172
+ assert(v._1 == column.getByte(v._2), " VectorType=" + column.getClass.getSimpleName)
181
173
}
182
174
}
183
175
184
176
testVector(" Short APIs" , 1024 , ShortType ) {
185
- ( column, memMode) =>
177
+ column =>
186
178
val seed = System .currentTimeMillis()
187
179
val random = new Random (seed)
188
180
val reference = mutable.ArrayBuffer .empty[Short ]
@@ -248,16 +240,13 @@ class ColumnarBatchSuite extends SparkFunSuite {
248
240
}
249
241
250
242
reference.zipWithIndex.foreach { v =>
251
- assert(v._1 == column.getShort(v._2), " Seed = " + seed + " Mem Mode=" + memMode)
252
- if (memMode == MemoryMode .OFF_HEAP ) {
253
- val addr = column.valuesNativeAddress()
254
- assert(v._1 == Platform .getShort(null , addr + 2 * v._2))
255
- }
243
+ assert(v._1 == column.getShort(v._2),
244
+ " Seed = " + seed + " VectorType=" + column.getClass.getSimpleName)
256
245
}
257
246
}
258
247
259
248
testVector(" Int APIs" , 1024 , IntegerType ) {
260
- ( column, memMode) =>
249
+ column =>
261
250
val seed = System .currentTimeMillis()
262
251
val random = new Random (seed)
263
252
val reference = mutable.ArrayBuffer .empty[Int ]
@@ -329,16 +318,13 @@ class ColumnarBatchSuite extends SparkFunSuite {
329
318
}
330
319
331
320
reference.zipWithIndex.foreach { v =>
332
- assert(v._1 == column.getInt(v._2), " Seed = " + seed + " Mem Mode=" + memMode)
333
- if (memMode == MemoryMode .OFF_HEAP ) {
334
- val addr = column.valuesNativeAddress()
335
- assert(v._1 == Platform .getInt(null , addr + 4 * v._2))
336
- }
321
+ assert(v._1 == column.getInt(v._2),
322
+ " Seed = " + seed + " VectorType=" + column.getClass.getSimpleName)
337
323
}
338
324
}
339
325
340
326
testVector(" Long APIs" , 1024 , LongType ) {
341
- ( column, memMode) =>
327
+ column =>
342
328
val seed = System .currentTimeMillis()
343
329
val random = new Random (seed)
344
330
val reference = mutable.ArrayBuffer .empty[Long ]
@@ -413,16 +399,12 @@ class ColumnarBatchSuite extends SparkFunSuite {
413
399
414
400
reference.zipWithIndex.foreach { v =>
415
401
assert(v._1 == column.getLong(v._2), " idx=" + v._2 +
416
- " Seed = " + seed + " MemMode=" + memMode)
417
- if (memMode == MemoryMode .OFF_HEAP ) {
418
- val addr = column.valuesNativeAddress()
419
- assert(v._1 == Platform .getLong(null , addr + 8 * v._2))
420
- }
402
+ " Seed = " + seed + " VectorType=" + column.getClass.getSimpleName)
421
403
}
422
404
}
423
405
424
406
testVector(" Float APIs" , 1024 , FloatType ) {
425
- ( column, memMode) =>
407
+ column =>
426
408
val seed = System .currentTimeMillis()
427
409
val random = new Random (seed)
428
410
val reference = mutable.ArrayBuffer .empty[Float ]
@@ -500,16 +482,13 @@ class ColumnarBatchSuite extends SparkFunSuite {
500
482
}
501
483
502
484
reference.zipWithIndex.foreach { v =>
503
- assert(v._1 == column.getFloat(v._2), " Seed = " + seed + " MemMode=" + memMode)
504
- if (memMode == MemoryMode .OFF_HEAP ) {
505
- val addr = column.valuesNativeAddress()
506
- assert(v._1 == Platform .getFloat(null , addr + 4 * v._2))
507
- }
485
+ assert(v._1 == column.getFloat(v._2),
486
+ " Seed = " + seed + " VectorType=" + column.getClass.getSimpleName)
508
487
}
509
488
}
510
489
511
490
testVector(" Double APIs" , 1024 , DoubleType ) {
512
- ( column, memMode) =>
491
+ column =>
513
492
val seed = System .currentTimeMillis()
514
493
val random = new Random (seed)
515
494
val reference = mutable.ArrayBuffer .empty[Double ]
@@ -587,16 +566,13 @@ class ColumnarBatchSuite extends SparkFunSuite {
587
566
}
588
567
589
568
reference.zipWithIndex.foreach { v =>
590
- assert(v._1 == column.getDouble(v._2), " Seed = " + seed + " MemMode=" + memMode)
591
- if (memMode == MemoryMode .OFF_HEAP ) {
592
- val addr = column.valuesNativeAddress()
593
- assert(v._1 == Platform .getDouble(null , addr + 8 * v._2))
594
- }
569
+ assert(v._1 == column.getDouble(v._2),
570
+ " Seed = " + seed + " VectorType=" + column.getClass.getSimpleName)
595
571
}
596
572
}
597
573
598
574
testVector(" String APIs" , 6 , StringType ) {
599
- ( column, memMode) =>
575
+ column =>
600
576
val reference = mutable.ArrayBuffer .empty[String ]
601
577
602
578
assert(column.arrayData().elementsAppended == 0 )
@@ -643,17 +619,17 @@ class ColumnarBatchSuite extends SparkFunSuite {
643
619
assert(column.arrayData().elementsAppended == 17 + (s + s).length)
644
620
645
621
reference.zipWithIndex.foreach { v =>
646
- assert(v._1.length == column.getArrayLength(v._2), " MemoryMode =" + memMode)
647
- assert(v._1 == column.getUTF8String (v._2).toString,
648
- " MemoryMode " + memMode )
622
+ val errMsg = " VectorType =" + column.getClass.getSimpleName
623
+ assert(v._1.length == column.getArrayLength (v._2), errMsg)
624
+ assert(v._1 == column.getUTF8String(v._2).toString, errMsg )
649
625
}
650
626
651
627
column.reset()
652
628
assert(column.arrayData().elementsAppended == 0 )
653
629
}
654
630
655
631
testVector(" Int Array" , 10 , new ArrayType (IntegerType , true )) {
656
- ( column, _) =>
632
+ column =>
657
633
658
634
// Fill the underlying data with all the arrays back to back.
659
635
val data = column.arrayData()
@@ -763,7 +739,7 @@ class ColumnarBatchSuite extends SparkFunSuite {
763
739
testVector(
764
740
" Struct Column" ,
765
741
10 ,
766
- new StructType ().add(" int" , IntegerType ).add(" double" , DoubleType )) { ( column, _) =>
742
+ new StructType ().add(" int" , IntegerType ).add(" double" , DoubleType )) { column =>
767
743
val c1 = column.getChildColumn(0 )
768
744
val c2 = column.getChildColumn(1 )
769
745
assert(c1.dataType() == IntegerType )
@@ -789,7 +765,7 @@ class ColumnarBatchSuite extends SparkFunSuite {
789
765
}
790
766
791
767
testVector(" Nest Array in Array" , 10 , new ArrayType (new ArrayType (IntegerType , true ), true )) {
792
- ( column, _) =>
768
+ column =>
793
769
val childColumn = column.arrayData()
794
770
val data = column.arrayData().arrayData()
795
771
(0 until 6 ).foreach {
@@ -822,7 +798,7 @@ class ColumnarBatchSuite extends SparkFunSuite {
822
798
testVector(
823
799
" Nest Struct in Array" ,
824
800
10 ,
825
- new ArrayType (structType, true )) { ( column, _) =>
801
+ new ArrayType (structType, true )) { column =>
826
802
val data = column.arrayData()
827
803
val c0 = data.getChildColumn(0 )
828
804
val c1 = data.getChildColumn(1 )
@@ -851,7 +827,7 @@ class ColumnarBatchSuite extends SparkFunSuite {
851
827
10 ,
852
828
new StructType ()
853
829
.add(" int" , IntegerType )
854
- .add(" array" , new ArrayType (IntegerType , true ))) { ( column, _) =>
830
+ .add(" array" , new ArrayType (IntegerType , true ))) { column =>
855
831
val c0 = column.getChildColumn(0 )
856
832
val c1 = column.getChildColumn(1 )
857
833
c0.putInt(0 , 0 )
@@ -880,7 +856,7 @@ class ColumnarBatchSuite extends SparkFunSuite {
880
856
testVector(
881
857
" Nest Struct in Struct" ,
882
858
10 ,
883
- new StructType ().add(" int" , IntegerType ).add(" struct" , subSchema)) { ( column, _) =>
859
+ new StructType ().add(" int" , IntegerType ).add(" struct" , subSchema)) { column =>
884
860
val c0 = column.getChildColumn(0 )
885
861
val c1 = column.getChildColumn(1 )
886
862
c0.putInt(0 , 0 )
0 commit comments