Skip to content

Commit

Permalink
Update
Browse files Browse the repository at this point in the history
  • Loading branch information
Chong Gao committed Jan 9, 2024
1 parent a28e6bc commit 251d18c
Showing 1 changed file with 3 additions and 19 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

package com.nvidia.spark.rapids

import org.apache.spark.sql.rapids.execution.TrampolineUtil
import org.apache.spark.sql.types._

object DataTypeUtils {
Expand All @@ -34,24 +35,7 @@ object DataTypeUtils {
* @return if contains date type.
*/
def hasDateOrTimestampType(t: DataType): Boolean = {
hasType(t, t => t.isInstanceOf[DateType] || t.isInstanceOf[TimestampType])
}

/**
* If the specified date type or its children have a true predicate
*
* @param t input data type.
* @param predicate predicate for a date type.
* @return true if date type or its children have a true predicate.
*/
private def hasType(t: DataType, predicate: DataType => Boolean): Boolean = {
t match {
case _ if predicate(t) => true
case MapType(keyType, valueType, _) =>
hasType(keyType, predicate) || hasType(valueType, predicate)
case ArrayType(elementType, _) => hasType(elementType, predicate)
case StructType(fields) => fields.exists(f => hasType(f.dataType, predicate))
case _ => false
}
TrampolineUtil.dataTypeExistsRecursively(t, e =>
e.isInstanceOf[DateType] || e.isInstanceOf[TimestampType])
}
}

0 comments on commit 251d18c

Please sign in to comment.