Skip to content

[SPARK-48019] Fix incorrect behavior in ColumnVector/ColumnarArray with dictionary and nulls #46254

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
gene-db wants to merge 3 commits into from
Closed

[SPARK-48019] Fix incorrect behavior in ColumnVector/ColumnarArray with dictionary and nulls #46254

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f66452b
Fix array index OOB ex for vectors with dictionary and nulls
gene-db Apr 26, 2024
4a163f1
Add copy() test (doesn't pass)
gene-db Apr 26, 2024
58572bf
avoid using UnsafeArrayData if nulls exist
gene-db Apr 26, 2024
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
5 changes: 5 additions & 0 deletions sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ColumnarArray.java
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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