diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/collectionOperations.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/collectionOperations.scala index 9253982f328..f378c1f2c14 100644 --- a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/collectionOperations.scala +++ b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/collectionOperations.scala @@ -162,15 +162,23 @@ case class GpuSlice(x: Expression, start: Expression, length: Expression) override def doColumnar(listS: GpuScalar, startCol: GpuColumnVector, lengthCol: GpuColumnVector): ColumnVector = { val numRows = startCol.getRowCount.toInt - withResource(GpuColumnVector.from(listS, numRows, dataType)) { listCol => - doColumnar(listCol, startCol, lengthCol) + // When list is null, return all nulls like the CPU does. + if (!listS.isValid) { + GpuColumnVector.columnVectorFromNull(numRows, dataType) + } else { + withResource(GpuColumnVector.from(listS, numRows, dataType)) { listCol => + doColumnar(listCol, startCol, lengthCol) + } } } override def doColumnar(listCol: GpuColumnVector, startS: GpuScalar, lengthS: GpuScalar): ColumnVector = { - // When either start or length is null, return all nulls like the CPU does. - if (listCol.getRowCount == listCol.numNulls() || !startS.isValid || !lengthS.isValid) { + // When input column is all nulls or either start or length is null, + // return all nulls like the CPU does. This matches CPU behavior for + // slice function with null inputs. + if (listCol.getRowCount == listCol.getBase.getNullCount || + !startS.isValid || !lengthS.isValid) { GpuColumnVector.columnVectorFromNull(listCol.getRowCount.toInt, dataType) } else { val list = listCol.getBase @@ -189,7 +197,7 @@ case class GpuSlice(x: Expression, start: Expression, length: Expression) override def doColumnar(listCol: GpuColumnVector, startS: GpuScalar, lengthCol: GpuColumnVector): ColumnVector = { // When start is null, return all nulls like the CPU does. - if (listCol.getRowCount == listCol.numNulls() || !startS.isValid) { + if (listCol.getRowCount == listCol.getBase.getNullCount || !startS.isValid) { GpuColumnVector.columnVectorFromNull(listCol.getRowCount.toInt, dataType) } else { val list = listCol.getBase @@ -211,11 +219,16 @@ case class GpuSlice(x: Expression, start: Expression, length: Expression) override def doColumnar(listCol: GpuColumnVector, startCol: GpuColumnVector, lengthS: GpuScalar): ColumnVector = { // When length is null, return all nulls like the CPU does. - if (listCol.getRowCount == listCol.numNulls() || !lengthS.isValid) { + if (listCol.getRowCount == listCol.getBase.getNullCount || !lengthS.isValid) { GpuColumnVector.columnVectorFromNull(listCol.getRowCount.toInt, dataType) } else { val list = listCol.getBase val start = startCol.getBase + withResource(GpuDivModLike.mergeNulls(list, start)) { mergedNulls => + if (mergedNulls.getRowCount == mergedNulls.getNullCount) { + return GpuColumnVector.columnVectorFromNull(listCol.getRowCount.toInt, dataType) + } + } val length = lengthS.getValue.asInstanceOf[Int] withResource(Scalar.fromInt(0)) { zero => if (start.contains(zero)) { @@ -231,12 +244,20 @@ case class GpuSlice(x: Expression, start: Expression, length: Expression) override def doColumnar(listCol: GpuColumnVector, startCol: GpuColumnVector, lengthCol: GpuColumnVector): ColumnVector = { - if (listCol.getRowCount == listCol.numNulls()) { + if (listCol.getRowCount == listCol.getBase.getNullCount) { GpuColumnVector.columnVectorFromNull(listCol.getRowCount.toInt, dataType) } else { val list = listCol.getBase val start = startCol.getBase val length = lengthCol.getBase + val mergedNullsList = withResource(GpuDivModLike.mergeNulls(list, start)) { mergedNulls => + GpuDivModLike.mergeNulls(mergedNulls, length) + } + withResource(mergedNullsList) { _ => + if (mergedNullsList.getRowCount == mergedNullsList.getNullCount) { + return GpuColumnVector.columnVectorFromNull(listCol.getRowCount.toInt, dataType) + } + } withResource(Scalar.fromInt(0)) { zero => if (start.contains(zero)) { throw RapidsErrorUtils.unexpectedValueForStartInFunctionError(prettyName)