Skip to content

Commit

Permalink
Merge pull request NVIDIA#2189 from NVIDIA/branch-0.5
Browse files Browse the repository at this point in the history 
[auto-merge] branch-0.5 to branch-0.6 [skip ci] [bot]
  • Loading branch information
@nvauto
nvauto authored Apr 19, 2021
2 parents 36d08f0 + 435a7b8 commit 37885f0
Showing 1 changed file with 75 additions and 129 deletions.
204 changes: 75 additions & 129 deletions sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuExpressions.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -114,22 +114,14 @@ abstract class GpuUnaryExpression extends UnaryExpression with GpuExpression {
}

override def columnarEval(batch: ColumnarBatch): Any = {
val input = child.columnarEval(batch)
try {
input match {
case vec: GpuColumnVector =>
doItColumnar(vec)
case other =>
withResource(GpuScalar.from(other, child.dataType)) { s =>
withResource(GpuColumnVector.from(s, batch.numRows(), child.dataType)) { vec =>
doItColumnar(vec)
}
withResourceIfAllowed(child.columnarEval(batch)) {
case vec: GpuColumnVector => doItColumnar(vec)
case other =>
withResource(GpuScalar.from(other, child.dataType)) { s =>
withResource(GpuColumnVector.from(s, batch.numRows(), child.dataType)) { vec =>
doItColumnar(vec)
}
}
} finally {
if (input.isInstanceOf[AutoCloseable]) {
input.asInstanceOf[AutoCloseable].close()
}
}
}
}
}
Expand All @@ -148,37 +140,27 @@ trait GpuBinaryExpression extends BinaryExpression with GpuExpression {
def doColumnar(numRows: Int, lhs: Scalar, rhs: Scalar): ColumnVector

override def columnarEval(batch: ColumnarBatch): Any = {
var lhs: Any = null
var rhs: Any = null
try {
lhs = left.columnarEval(batch)
rhs = right.columnarEval(batch)

(lhs, rhs) match {
case (l: GpuColumnVector, r: GpuColumnVector) =>
GpuColumnVector.from(doColumnar(l, r), dataType)
case (l, r: GpuColumnVector) =>
withResource(GpuScalar.from(l, left.dataType)) { scalar =>
GpuColumnVector.from(doColumnar(scalar, r), dataType)
}
case (l: GpuColumnVector, r) =>
withResource(GpuScalar.from(r, right.dataType)) { scalar =>
GpuColumnVector.from(doColumnar(l, scalar), dataType)
}
case (l, r) if l != null && r != null =>
withResource(GpuScalar.from(l, left.dataType)) { leftScalar =>
withResource(GpuScalar.from(r, right.dataType)) { rightScalar =>
GpuColumnVector.from(doColumnar(batch.numRows(), leftScalar, rightScalar), dataType)
withResourceIfAllowed(left.columnarEval(batch)) { lhs =>
withResourceIfAllowed(right.columnarEval(batch)) { rhs =>
(lhs, rhs) match {
case (l: GpuColumnVector, r: GpuColumnVector) =>
GpuColumnVector.from(doColumnar(l, r), dataType)
case (l, r: GpuColumnVector) =>
withResource(GpuScalar.from(l, left.dataType)) { scalar =>
GpuColumnVector.from(doColumnar(scalar, r), dataType)
}
}
case _ => null
}
} finally {
if (lhs.isInstanceOf[AutoCloseable]) {
lhs.asInstanceOf[AutoCloseable].close()
}
if (rhs.isInstanceOf[AutoCloseable]) {
rhs.asInstanceOf[AutoCloseable].close()
case (l: GpuColumnVector, r) =>
withResource(GpuScalar.from(r, right.dataType)) { scalar =>
GpuColumnVector.from(doColumnar(l, scalar), dataType)
}
case (l, r) if l != null && r != null =>
withResource(GpuScalar.from(l, left.dataType)) { leftScalar =>
withResource(GpuScalar.from(r, right.dataType)) { rightScalar =>
GpuColumnVector.from(doColumnar(batch.numRows(), leftScalar, rightScalar), dataType)
}
}
case _ => null
}
}
}
}
Expand Down Expand Up @@ -258,14 +240,13 @@ trait GpuString2TrimExpression extends String2TrimExpression with GpuExpression

override def columnarEval(batch: ColumnarBatch): Any = {
val trim = GpuExpressionsUtils.getTrimString(trimStr)
val shouldBeColumn = srcStr.columnarEval(batch)
try {
withResourceIfAllowed(srcStr.columnarEval(batch)) { shouldBeColumn =>
// We know the first parameter is not a Literal, because trim(Literal, Literal) would already
// have been optimized out
val column = shouldBeColumn.asInstanceOf[GpuColumnVector]
if (trim == null) {
withResource(GpuScalar.from(null, StringType)) { nullScalar =>
GpuColumnVector.from(nullScalar, column.getRowCount().toInt, StringType)
GpuColumnVector.from(nullScalar, column.getRowCount.toInt, StringType)
}
} else if (trim.isEmpty) {
column.incRefCount() // This is a noop
Expand All @@ -274,10 +255,6 @@ trait GpuString2TrimExpression extends String2TrimExpression with GpuExpression
strippedColumnVector(column, t)
}
}
} finally {
if (shouldBeColumn.isInstanceOf[AutoCloseable]) {
shouldBeColumn.asInstanceOf[AutoCloseable].close()
}
}
}
}
Expand All @@ -295,85 +272,54 @@ trait GpuTernaryExpression extends TernaryExpression with GpuExpression {
def doColumnar(numRows: Int, val0: Scalar, val1: Scalar, val2: Scalar): ColumnVector

override def columnarEval(batch: ColumnarBatch): Any = {
var val0: Any = null
var val1: Any = null
var val2: Any = null
try {
val0 = children(0).columnarEval(batch)
val1 = children(1).columnarEval(batch)
val2 = children(2).columnarEval(batch)

(val0, val1, val2) match {
case (v0: GpuColumnVector, v1: GpuColumnVector, v2: GpuColumnVector) =>
doColumnar(v0, v1, v2)
case (v0, v1: GpuColumnVector, v2: GpuColumnVector) =>
val scalar0 = GpuScalar.from(v0, children(0).dataType)
try {
GpuColumnVector.from(doColumnar(scalar0, v1, v2), dataType)
} finally {
scalar0.close()
}
case (v0: GpuColumnVector, v1, v2: GpuColumnVector) =>
val scalar1 = GpuScalar.from(v1, children(1).dataType)
try {
GpuColumnVector.from(doColumnar(v0, scalar1, v2), dataType)
} finally {
scalar1.close()
}
case (v0: GpuColumnVector, v1: GpuColumnVector, v2) =>
val scalar2 = GpuScalar.from(v2, children(2).dataType)
try {
GpuColumnVector.from(doColumnar(v0, v1, scalar2), dataType)
} finally {
scalar2.close()
}
case (v0, v1, v2: GpuColumnVector) =>
val scalar0 = GpuScalar.from(v0, children(0).dataType)
val scalar1 = GpuScalar.from(v1, children(1).dataType)
try {
GpuColumnVector.from(doColumnar(scalar0, scalar1, v2), dataType)
} finally {
scalar0.close()
scalar1.close()
}
case (v0, v1: GpuColumnVector, v2) =>
val scalar0 = GpuScalar.from(v0, children(0).dataType)
val scalar2 = GpuScalar.from(v2, children(2).dataType)
try {
GpuColumnVector.from(doColumnar(scalar0, v1, scalar2), dataType)
} finally {
scalar0.close()
scalar2.close()
}
case (v0: GpuColumnVector, v1, v2) =>
val scalar1 = GpuScalar.from(v1, children(1).dataType)
val scalar2 = GpuScalar.from(v2, children(2).dataType)
try {
GpuColumnVector.from(doColumnar(v0, scalar1, scalar2), dataType)
} finally {
scalar1.close()
scalar2.close()
}
case (v0, v1, v2) if v0 != null && v1 != null && v2 != null =>
withResource(GpuScalar.from(v0, children(0).dataType)) { v0Scalar =>
withResource(GpuScalar.from(v1, children(1).dataType)) { v1Scalar =>
withResource(GpuScalar.from(v2, children(2).dataType)) { v2Scalar =>
GpuColumnVector.from(doColumnar(batch.numRows(), v0Scalar, v1Scalar, v2Scalar),
dataType)
withResourceIfAllowed(children(0).columnarEval(batch)) { val0 =>
withResourceIfAllowed(children(1).columnarEval(batch)) { val1 =>
withResourceIfAllowed(children(2).columnarEval(batch)) { val2 =>
(val0, val1, val2) match {
case (v0: GpuColumnVector, v1: GpuColumnVector, v2: GpuColumnVector) =>
doColumnar(v0, v1, v2)
case (v0, v1: GpuColumnVector, v2: GpuColumnVector) =>
withResource(GpuScalar.from(v0, children(0).dataType)) { scalar0 =>
GpuColumnVector.from(doColumnar(scalar0, v1, v2), dataType)
}
}
case (v0: GpuColumnVector, v1, v2: GpuColumnVector) =>
withResource(GpuScalar.from(v1, children(1).dataType)) { scalar1 =>
GpuColumnVector.from(doColumnar(v0, scalar1, v2), dataType)
}
case (v0: GpuColumnVector, v1: GpuColumnVector, v2) =>
withResource(GpuScalar.from(v2, children(2).dataType)) { scalar2 =>
GpuColumnVector.from(doColumnar(v0, v1, scalar2), dataType)
}
case (v0, v1, v2: GpuColumnVector) =>
withResource(GpuScalar.from(v0, children(0).dataType)) { scalar0 =>
withResource(GpuScalar.from(v1, children(1).dataType)) { scalar1 =>
GpuColumnVector.from(doColumnar(scalar0, scalar1, v2), dataType)
}
}
case (v0, v1: GpuColumnVector, v2) =>
withResource(GpuScalar.from(v0, children(0).dataType)) { scalar0 =>
withResource(GpuScalar.from(v2, children(2).dataType)) { scalar2 =>
GpuColumnVector.from(doColumnar(scalar0, v1, scalar2), dataType)
}
}
case (v0: GpuColumnVector, v1, v2) =>
withResource(GpuScalar.from(v1, children(1).dataType)) { scalar1 =>
withResource(GpuScalar.from(v2, children(2).dataType)) { scalar2 =>
GpuColumnVector.from(doColumnar(v0, scalar1, scalar2), dataType)
}
}
case (v0, v1, v2) if v0 != null && v1 != null && v2 != null =>
withResource(GpuScalar.from(v0, children(0).dataType)) { v0Scalar =>
withResource(GpuScalar.from(v1, children(1).dataType)) { v1Scalar =>
withResource(GpuScalar.from(v2, children(2).dataType)) { v2Scalar =>
GpuColumnVector.from(doColumnar(batch.numRows(), v0Scalar, v1Scalar, v2Scalar),
dataType)
}
}
}
case _ => null
}
case _ => null
}
} finally {
if (val0.isInstanceOf[AutoCloseable]) {
val0.asInstanceOf[AutoCloseable].close()
}
if (val1.isInstanceOf[AutoCloseable]) {
val1.asInstanceOf[AutoCloseable].close()
}
if (val2.isInstanceOf[AutoCloseable]) {
val2.asInstanceOf[AutoCloseable].close()
}
}
}
}
Expand Down

0 comments on commit 37885f0

Please sign in to comment.