Skip to content

Commit

Permalink
[SPARK-48019] Fix incorrect behavior in ColumnVector/ColumnarArray wi…
Browse files Browse the repository at this point in the history 
…th dictionary and nulls

### What changes were proposed in this pull request?

This fixes how `ColumnVector` handles copying arrays when the vector has a dictionary and null values. The possible issues with the previous implementation:
- An `ArrayIndexOutOfBoundsException` may be thrown when the `ColumnVector` has nulls and dictionaries. This is because the dictionary id for `null` entries might be invalid and should not be used for `null` entries.
- Copying a `ColumnarArray` (which contains a `ColumnVector`) is incorrect, if it contains `null` entries. This is because copying a primitive array does not take into account the `null` entries, so all the null entries get lost.

### Why are the changes needed?

These changes are needed to avoid `ArrayIndexOutOfBoundsException` and to produce correct results when copying `ColumnarArray`.

### Does this PR introduce _any_ user-facing change?

The only user facing changes are to fix existing errors and incorrect results.

### How was this patch tested?

Added new unit tests.

### Was this patch authored or co-authored using generative AI tooling?

No.

Closes apache#46254 from gene-db/dictionary-nulls.

Authored-by: Gene Pang <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
  • Loading branch information
@gene-db @cloud-fan
gene-db authored and cloud-fan committed Apr 28, 2024
1 parent 64d3219 commit 76ce6b0
Show file tree
Hide file tree
Showing 4 changed files with 215 additions and 12 deletions.
5 changes: 5 additions & 0 deletions sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ColumnarArray.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,11 @@ public int numElements() {
public ArrayData copy() {
DataType dt = data.dataType();

if (data.hasNull()) {
// UnsafeArrayData cannot be used if there are any nulls.
return new GenericArrayData(toObjectArray(dt)).copy();
}

if (dt instanceof BooleanType) {
return UnsafeArrayData.fromPrimitiveArray(toBooleanArray());
} else if (dt instanceof ByteType) {
Expand Down
24 changes: 18 additions & 6 deletions sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/OffHeapColumnVector.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -218,7 +218,9 @@ public byte[] getBytes(int rowId, int count) {
Platform.copyMemory(null, data + rowId, array, Platform.BYTE_ARRAY_OFFSET, count);
} else {
for (int i = 0; i < count; i++) {
array[i] = (byte) dictionary.decodeToInt(dictionaryIds.getDictId(rowId + i));
if (!isNullAt(rowId + i)) {
array[i] = (byte) dictionary.decodeToInt(dictionaryIds.getDictId(rowId + i));
}
}
}
return array;
Expand Down Expand Up @@ -279,7 +281,9 @@ public short[] getShorts(int rowId, int count) {
Platform.copyMemory(null, data + rowId * 2L, array, Platform.SHORT_ARRAY_OFFSET, count * 2L);
} else {
for (int i = 0; i < count; i++) {
array[i] = (short) dictionary.decodeToInt(dictionaryIds.getDictId(rowId + i));
if (!isNullAt(rowId + i)) {
array[i] = (short) dictionary.decodeToInt(dictionaryIds.getDictId(rowId + i));
}
}
}
return array;
Expand Down Expand Up @@ -345,7 +349,9 @@ public int[] getInts(int rowId, int count) {
Platform.copyMemory(null, data + rowId * 4L, array, Platform.INT_ARRAY_OFFSET, count * 4L);
} else {
for (int i = 0; i < count; i++) {
array[i] = dictionary.decodeToInt(dictionaryIds.getDictId(rowId + i));
if (!isNullAt(rowId + i)) {
array[i] = dictionary.decodeToInt(dictionaryIds.getDictId(rowId + i));
}
}
}
return array;
Expand Down Expand Up @@ -423,7 +429,9 @@ public long[] getLongs(int rowId, int count) {
Platform.copyMemory(null, data + rowId * 8L, array, Platform.LONG_ARRAY_OFFSET, count * 8L);
} else {
for (int i = 0; i < count; i++) {
array[i] = dictionary.decodeToLong(dictionaryIds.getDictId(rowId + i));
if (!isNullAt(rowId + i)) {
array[i] = dictionary.decodeToLong(dictionaryIds.getDictId(rowId + i));
}
}
}
return array;
Expand Down Expand Up @@ -487,7 +495,9 @@ public float[] getFloats(int rowId, int count) {
Platform.copyMemory(null, data + rowId * 4L, array, Platform.FLOAT_ARRAY_OFFSET, count * 4L);
} else {
for (int i = 0; i < count; i++) {
array[i] = dictionary.decodeToFloat(dictionaryIds.getDictId(rowId + i));
if (!isNullAt(rowId + i)) {
array[i] = dictionary.decodeToFloat(dictionaryIds.getDictId(rowId + i));
}
}
}
return array;
Expand Down Expand Up @@ -553,7 +563,9 @@ public double[] getDoubles(int rowId, int count) {
count * 8L);
} else {
for (int i = 0; i < count; i++) {
array[i] = dictionary.decodeToDouble(dictionaryIds.getDictId(rowId + i));
if (!isNullAt(rowId + i)) {
array[i] = dictionary.decodeToDouble(dictionaryIds.getDictId(rowId + i));
}
}
}
return array;
Expand Down
24 changes: 18 additions & 6 deletions sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/OnHeapColumnVector.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -216,7 +216,9 @@ public byte[] getBytes(int rowId, int count) {
System.arraycopy(byteData, rowId, array, 0, count);
} else {
for (int i = 0; i < count; i++) {
array[i] = (byte) dictionary.decodeToInt(dictionaryIds.getDictId(rowId + i));
if (!isNullAt(rowId + i)) {
array[i] = (byte) dictionary.decodeToInt(dictionaryIds.getDictId(rowId + i));
}
}
}
return array;
Expand Down Expand Up @@ -276,7 +278,9 @@ public short[] getShorts(int rowId, int count) {
System.arraycopy(shortData, rowId, array, 0, count);
} else {
for (int i = 0; i < count; i++) {
array[i] = (short) dictionary.decodeToInt(dictionaryIds.getDictId(rowId + i));
if (!isNullAt(rowId + i)) {
array[i] = (short) dictionary.decodeToInt(dictionaryIds.getDictId(rowId + i));
}
}
}
return array;
Expand Down Expand Up @@ -337,7 +341,9 @@ public int[] getInts(int rowId, int count) {
System.arraycopy(intData, rowId, array, 0, count);
} else {
for (int i = 0; i < count; i++) {
array[i] = dictionary.decodeToInt(dictionaryIds.getDictId(rowId + i));
if (!isNullAt(rowId + i)) {
array[i] = dictionary.decodeToInt(dictionaryIds.getDictId(rowId + i));
}
}
}
return array;
Expand Down Expand Up @@ -409,7 +415,9 @@ public long[] getLongs(int rowId, int count) {
System.arraycopy(longData, rowId, array, 0, count);
} else {
for (int i = 0; i < count; i++) {
array[i] = dictionary.decodeToLong(dictionaryIds.getDictId(rowId + i));
if (!isNullAt(rowId + i)) {
array[i] = dictionary.decodeToLong(dictionaryIds.getDictId(rowId + i));
}
}
}
return array;
Expand Down Expand Up @@ -466,7 +474,9 @@ public float[] getFloats(int rowId, int count) {
System.arraycopy(floatData, rowId, array, 0, count);
} else {
for (int i = 0; i < count; i++) {
array[i] = dictionary.decodeToFloat(dictionaryIds.getDictId(rowId + i));
if (!isNullAt(rowId + i)) {
array[i] = dictionary.decodeToFloat(dictionaryIds.getDictId(rowId + i));
}
}
}
return array;
Expand Down Expand Up @@ -525,7 +535,9 @@ public double[] getDoubles(int rowId, int count) {
System.arraycopy(doubleData, rowId, array, 0, count);
} else {
for (int i = 0; i < count; i++) {
array[i] = dictionary.decodeToDouble(dictionaryIds.getDictId(rowId + i));
if (!isNullAt(rowId + i)) {
array[i] = dictionary.decodeToDouble(dictionaryIds.getDictId(rowId + i));
}
}
}
return array;
Expand Down
174 changes: 174 additions & 0 deletions sql/core/src/test/scala/org/apache/spark/sql/execution/vectorized/ColumnVectorSuite.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -476,6 +476,180 @@ class ColumnVectorSuite extends SparkFunSuite with SQLHelper {
assert(testVector.getDoubles(0, 3)(2) == 1342.17729d)
}

def check(expected: Seq[Any], testVector: WritableColumnVector): Unit = {
expected.zipWithIndex.foreach {
case (v: Integer, idx) =>
assert(testVector.getInt(idx) == v)
assert(testVector.getInts(0, testVector.capacity)(idx) == v)
case (v: Short, idx) =>
assert(testVector.getShort(idx) == v)
assert(testVector.getShorts(0, testVector.capacity)(idx) == v)
case (v: Byte, idx) =>
assert(testVector.getByte(idx) == v)
assert(testVector.getBytes(0, testVector.capacity)(idx) == v)
case (v: Long, idx) =>
assert(testVector.getLong(idx) == v)
assert(testVector.getLongs(0, testVector.capacity)(idx) == v)
case (v: Float, idx) =>
assert(testVector.getFloat(idx) == v)
assert(testVector.getFloats(0, testVector.capacity)(idx) == v)
case (v: Double, idx) =>
assert(testVector.getDouble(idx) == v)
assert(testVector.getDoubles(0, testVector.capacity)(idx) == v)
case (null, idx) => testVector.isNullAt(idx)
case (_, idx) => assert(false, s"Unexpected value at $idx")
}

// Verify ColumnarArray.copy() works as expected
val arr = new ColumnarArray(testVector, 0, testVector.capacity)
assert(arr.toSeq(testVector.dataType) == expected)
assert(arr.copy().toSeq(testVector.dataType) == expected)
}

testVectors("getInts with dictionary and nulls", 3, IntegerType) { testVector =>
// Validate without dictionary
val expected = Seq(1, null, 3)
expected.foreach {
case i: Integer => testVector.appendInt(i)
case _ => testVector.appendNull()
}
check(expected, testVector)

// Validate with dictionary
val expectedDictionary = Seq(7, null, 9)
val dictArray = (Seq(-1, -1) ++ expectedDictionary.map {
case i: Integer => i.toInt
case _ => -1
}).toArray
val dict = new ColumnDictionary(dictArray)
testVector.setDictionary(dict)
testVector.reserveDictionaryIds(3)
testVector.getDictionaryIds.putInt(0, 2)
testVector.getDictionaryIds.putInt(1, -1) // This is a null, so the entry should be ignored
testVector.getDictionaryIds.putInt(2, 4)
check(expectedDictionary, testVector)
}

testVectors("getShorts with dictionary and nulls", 3, ShortType) { testVector =>
// Validate without dictionary
val expected = Seq(1.toShort, null, 3.toShort)
expected.foreach {
case i: Short => testVector.appendShort(i)
case _ => testVector.appendNull()
}
check(expected, testVector)

// Validate with dictionary
val expectedDictionary = Seq(7.toShort, null, 9.toShort)
val dictArray = (Seq(-1, -1) ++ expectedDictionary.map {
case i: Short => i.toInt
case _ => -1
}).toArray
val dict = new ColumnDictionary(dictArray)
testVector.setDictionary(dict)
testVector.reserveDictionaryIds(3)
testVector.getDictionaryIds.putInt(0, 2)
testVector.getDictionaryIds.putInt(1, -1) // This is a null, so the entry should be ignored
testVector.getDictionaryIds.putInt(2, 4)
check(expectedDictionary, testVector)
}

testVectors("getBytes with dictionary and nulls", 3, ByteType) { testVector =>
// Validate without dictionary
val expected = Seq(1.toByte, null, 3.toByte)
expected.foreach {
case i: Byte => testVector.appendByte(i)
case _ => testVector.appendNull()
}
check(expected, testVector)

// Validate with dictionary
val expectedDictionary = Seq(7.toByte, null, 9.toByte)
val dictArray = (Seq(-1, -1) ++ expectedDictionary.map {
case i: Byte => i.toInt
case _ => -1
}).toArray
val dict = new ColumnDictionary(dictArray)
testVector.setDictionary(dict)
testVector.reserveDictionaryIds(3)
testVector.getDictionaryIds.putInt(0, 2)
testVector.getDictionaryIds.putInt(1, -1) // This is a null, so the entry should be ignored
testVector.getDictionaryIds.putInt(2, 4)
check(expectedDictionary, testVector)
}

testVectors("getLongs with dictionary and nulls", 3, LongType) { testVector =>
// Validate without dictionary
val expected = Seq(2147483L, null, 2147485L)
expected.foreach {
case i: Long => testVector.appendLong(i)
case _ => testVector.appendNull()
}
check(expected, testVector)

// Validate with dictionary
val expectedDictionary = Seq(2147483648L, null, 2147483650L)
val dictArray = (Seq(-1L, -1L) ++ expectedDictionary.map {
case i: Long => i
case _ => -1L
}).toArray
val dict = new ColumnDictionary(dictArray)
testVector.setDictionary(dict)
testVector.reserveDictionaryIds(3)
testVector.getDictionaryIds.putInt(0, 2)
testVector.getDictionaryIds.putInt(1, -1) // This is a null, so the entry should be ignored
testVector.getDictionaryIds.putInt(2, 4)
check(expectedDictionary, testVector)
}

testVectors("getFloats with dictionary and nulls", 3, FloatType) { testVector =>
// Validate without dictionary
val expected = Seq(1.1f, null, 3.3f)
expected.foreach {
case i: Float => testVector.appendFloat(i)
case _ => testVector.appendNull()
}
check(expected, testVector)

// Validate with dictionary
val expectedDictionary = Seq(0.1f, null, 0.3f)
val dictArray = (Seq(-1f, -1f) ++ expectedDictionary.map {
case i: Float => i
case _ => -1f
}).toArray
val dict = new ColumnDictionary(dictArray)
testVector.setDictionary(dict)
testVector.reserveDictionaryIds(3)
testVector.getDictionaryIds.putInt(0, 2)
testVector.getDictionaryIds.putInt(1, -1) // This is a null, so the entry should be ignored
testVector.getDictionaryIds.putInt(2, 4)
check(expectedDictionary, testVector)
}

testVectors("getDoubles with dictionary and nulls", 3, DoubleType) { testVector =>
// Validate without dictionary
val expected = Seq(1.1d, null, 3.3d)
expected.foreach {
case i: Double => testVector.appendDouble(i)
case _ => testVector.appendNull()
}
check(expected, testVector)

// Validate with dictionary
val expectedDictionary = Seq(1342.17727d, null, 1342.17729d)
val dictArray = (Seq(-1d, -1d) ++ expectedDictionary.map {
case i: Double => i
case _ => -1d
}).toArray
val dict = new ColumnDictionary(dictArray)
testVector.setDictionary(dict)
testVector.reserveDictionaryIds(3)
testVector.getDictionaryIds.putInt(0, 2)
testVector.getDictionaryIds.putInt(1, -1) // This is a null, so the entry should be ignored
testVector.getDictionaryIds.putInt(2, 4)
check(expectedDictionary, testVector)
}

test("[SPARK-22092] off-heap column vector reallocation corrupts array data") {
withVector(new OffHeapColumnVector(8, arrayType)) { testVector =>
val data = testVector.arrayData()
Expand Down

0 comments on commit 76ce6b0

Please sign in to comment.