You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ag...@apache.org on 2025/05/16 06:00:22 UTC
(datafusion-comet) branch main updated: chore: Refactor DataTypeSupport (#1741)
This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new c89e4565d chore: Refactor DataTypeSupport (#1741)
c89e4565d is described below
commit c89e4565d8f381d82a0d9d5b6f4a2385fb284506
Author: Andy Grove <ag...@apache.org>
AuthorDate: Fri May 16 00:00:17 2025 -0600
chore: Refactor DataTypeSupport (#1741)
---
.../apache/comet/CometSparkSessionExtensions.scala | 7 +-
.../scala/org/apache/comet/DataTypeSupport.scala | 90 +++++++---------------
.../org/apache/comet/serde/QueryPlanSerde.scala | 4 +-
.../spark/sql/comet/CometBatchScanExec.scala | 13 +++-
.../spark/sql/comet/CometNativeScanExec.scala | 21 +++--
.../org/apache/spark/sql/comet/CometScanExec.scala | 41 +++++-----
.../spark/sql/comet/CometSparkToColumnarExec.scala | 6 +-
.../scala/org/apache/comet/CometCastSuite.scala | 2 +-
.../org/apache/comet/CometExpressionSuite.scala | 2 +-
.../comet/exec/CometColumnarShuffleSuite.scala | 2 +-
.../apache/comet/parquet/ParquetReadSuite.scala | 16 ++--
.../scala/org/apache/spark/sql/CometTestBase.scala | 2 +-
.../sql/comet/ParquetDatetimeRebaseSuite.scala | 12 +--
13 files changed, 96 insertions(+), 122 deletions(-)
diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
index 141429125..dce2d40a1 100644
--- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
+++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
@@ -242,10 +242,15 @@ object CometSparkSessionExtensions extends Logging {
org.apache.spark.SPARK_VERSION >= "4.0"
}
- def usingDataFusionParquetExec(conf: SQLConf): Boolean =
+ def usingDataSourceExec(conf: SQLConf): Boolean =
Seq(CometConf.SCAN_NATIVE_ICEBERG_COMPAT, CometConf.SCAN_NATIVE_DATAFUSION).contains(
CometConf.COMET_NATIVE_SCAN_IMPL.get(conf))
+ def usingDataSourceExecWithIncompatTypes(conf: SQLConf): Boolean = {
+ usingDataSourceExec(conf) &&
+ !CometConf.COMET_SCAN_ALLOW_INCOMPATIBLE.get(conf)
+ }
+
/**
* Whether we should override Spark memory configuration for Comet. This only returns true when
* Comet native execution is enabled and/or Comet shuffle is enabled and Comet doesn't use
diff --git a/spark/src/main/scala/org/apache/comet/DataTypeSupport.scala b/spark/src/main/scala/org/apache/comet/DataTypeSupport.scala
index 8b05c2038..917aa9569 100644
--- a/spark/src/main/scala/org/apache/comet/DataTypeSupport.scala
+++ b/spark/src/main/scala/org/apache/comet/DataTypeSupport.scala
@@ -21,79 +21,50 @@ package org.apache.comet
import scala.collection.mutable.ListBuffer
-import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types._
-import org.apache.comet.DataTypeSupport.{usingParquetExecWithIncompatTypes, ARRAY_ELEMENT, MAP_KEY, MAP_VALUE}
+import org.apache.comet.DataTypeSupport.{ARRAY_ELEMENT, MAP_KEY, MAP_VALUE}
trait DataTypeSupport {
/**
- * This can be overridden to provide support for specific data types that aren't universally
- * supported. This will only be called for non-universally supported types.
+ * Checks if this schema is supported by checking if each field in the schema is supported.
*
- * @param dt
- * the data type to check
+ * @param schema
+ * the schema to check the fields of
* @return
- * true if the datatype is supported
+ * true if all fields in the schema are supported
*/
- def isAdditionallySupported(
- dt: DataType,
- name: String,
- fallbackReasons: ListBuffer[String]): Boolean = false
-
- private def isGloballySupported(dt: DataType, fallbackReasons: ListBuffer[String]): Boolean =
- dt match {
- case ByteType | ShortType if usingParquetExecWithIncompatTypes(SQLConf.get) =>
- fallbackReasons += s"${CometConf.COMET_SCAN_ALLOW_INCOMPATIBLE.key} is false"
- false
- case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType |
- BinaryType | StringType | _: DecimalType | DateType | TimestampType =>
- true
- case t: DataType if t.typeName == "timestamp_ntz" =>
- true
- case other =>
- fallbackReasons += s"Unsupported type: $other"
- false
- }
+ def isSchemaSupported(schema: StructType, fallbackReasons: ListBuffer[String]): Boolean = {
+ schema.fields.forall(f => isTypeSupported(f.dataType, f.name, fallbackReasons))
+ }
/**
- * Checks if this schema is supported by checking if each field in the struct is supported.
- *
- * @param struct
- * the struct to check the fields of
- * @return
- * true if all fields in the struct are supported
+ * Determine if Comet supports a data type. This method can be overridden by specific operators
+ * as needed.
*/
- def isSchemaSupported(struct: StructType, fallbackReasons: ListBuffer[String]): Boolean = {
- struct.fields.forall(f => isTypeSupported(f.dataType, f.name, fallbackReasons))
- }
-
def isTypeSupported(
dt: DataType,
name: String,
fallbackReasons: ListBuffer[String]): Boolean = {
- if (isGloballySupported(dt, fallbackReasons) || isAdditionallySupported(
- dt,
- name,
- fallbackReasons)) {
- // If complex types are supported, we additionally want to recurse into their children
- dt match {
- case StructType(fields) =>
- fields.forall(f => isTypeSupported(f.dataType, f.name, fallbackReasons))
- case ArrayType(elementType, _) =>
- isTypeSupported(elementType, ARRAY_ELEMENT, fallbackReasons)
- case MapType(keyType, valueType, _) =>
- isTypeSupported(keyType, MAP_KEY, fallbackReasons) && isTypeSupported(
- valueType,
- MAP_VALUE,
- fallbackReasons)
- // Not a complex type
- case _ => true
- }
- } else {
- fallbackReasons += s"Unsupported ${name} of type ${dt}"
- false
+
+ dt match {
+ case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType |
+ BinaryType | StringType | _: DecimalType | DateType | TimestampType |
+ TimestampNTZType =>
+ true
+ case StructType(fields) =>
+ fields.forall(f => isTypeSupported(f.dataType, f.name, fallbackReasons))
+ case ArrayType(elementType, _) =>
+ isTypeSupported(elementType, ARRAY_ELEMENT, fallbackReasons)
+ case MapType(keyType, valueType, _) =>
+ isTypeSupported(keyType, MAP_KEY, fallbackReasons) && isTypeSupported(
+ valueType,
+ MAP_VALUE,
+ fallbackReasons)
+ case _ =>
+ fallbackReasons += s"Unsupported ${name} of type ${dt}"
+ false
}
}
}
@@ -102,9 +73,4 @@ object DataTypeSupport {
val ARRAY_ELEMENT = "array element"
val MAP_KEY = "map key"
val MAP_VALUE = "map value"
-
- def usingParquetExecWithIncompatTypes(conf: SQLConf): Boolean = {
- CometSparkSessionExtensions.usingDataFusionParquetExec(conf) &&
- !CometConf.COMET_SCAN_ALLOW_INCOMPATIBLE.get(conf)
- }
}
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
index 562a4ba86..f49ff2099 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -48,7 +48,7 @@ import org.apache.spark.sql.types._
import org.apache.spark.unsafe.types.UTF8String
import org.apache.comet.CometConf
-import org.apache.comet.CometSparkSessionExtensions.{isCometScan, usingDataFusionParquetExec, withInfo}
+import org.apache.comet.CometSparkSessionExtensions.{isCometScan, usingDataSourceExec, withInfo}
import org.apache.comet.expressions._
import org.apache.comet.serde.ExprOuterClass.{AggExpr, DataType => ProtoDataType, Expr, ScalarFunc}
import org.apache.comet.serde.ExprOuterClass.DataType._
@@ -2766,7 +2766,7 @@ object QueryPlanSerde extends Logging with CometExprShim {
// - Native datafusion reader enabled (experimental) OR
// - conversion from Parquet/JSON enabled
allowComplex =
- usingDataFusionParquetExec(conf) || CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED
+ usingDataSourceExec(conf) || CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED
.get(conf) || CometConf.COMET_CONVERT_FROM_JSON_ENABLED.get(conf))) =>
// These operators are source of Comet native execution chain
val scanBuilder = OperatorOuterClass.Scan.newBuilder()
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometBatchScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometBatchScanExec.scala
index 220df65ac..8c779e8dc 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometBatchScanExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometBatchScanExec.scala
@@ -19,6 +19,8 @@
package org.apache.spark.sql.comet
+import scala.collection.mutable.ListBuffer
+
import org.apache.spark.rdd._
import org.apache.spark.sql.catalyst._
import org.apache.spark.sql.catalyst.expressions.{Attribute, DynamicPruningExpression, Expression, Literal, SortOrder}
@@ -28,6 +30,7 @@ import org.apache.spark.sql.connector.read._
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.datasources.v2._
import org.apache.spark.sql.execution.metric._
+import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StructType}
import org.apache.spark.sql.vectorized._
import com.google.common.base.Objects
@@ -150,4 +153,12 @@ case class CometBatchScanExec(wrapped: BatchScanExec, runtimeFilters: Seq[Expres
override def supportsColumnar: Boolean = true
}
-object CometBatchScanExec extends DataTypeSupport
+object CometBatchScanExec extends DataTypeSupport {
+ override def isTypeSupported(
+ dt: DataType,
+ name: String,
+ fallbackReasons: ListBuffer[String]): Boolean = dt match {
+ case _: StructType | _: ArrayType | _: MapType => false
+ case _ => super.isTypeSupported(dt, name, fallbackReasons)
+ }
+}
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala
index 100b6d00e..cfbae8b94 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala
@@ -31,13 +31,14 @@ import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, UnknownPartit
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.datasources._
import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types._
import org.apache.spark.util.collection._
import com.google.common.base.Objects
-import org.apache.comet.DataTypeSupport
-import org.apache.comet.DataTypeSupport.{ARRAY_ELEMENT, MAP_KEY, MAP_VALUE}
+import org.apache.comet.{CometConf, DataTypeSupport}
+import org.apache.comet.CometSparkSessionExtensions.usingDataSourceExecWithIncompatTypes
import org.apache.comet.parquet.CometParquetFileFormat
import org.apache.comet.serde.OperatorOuterClass.Operator
@@ -231,20 +232,16 @@ object CometNativeScanExec extends DataTypeSupport {
batchScanExec
}
- override def isAdditionallySupported(
+ override def isTypeSupported(
dt: DataType,
name: String,
fallbackReasons: ListBuffer[String]): Boolean = {
dt match {
- case s: StructType =>
- s.fields.forall(f => isTypeSupported(f.dataType, f.name, fallbackReasons))
- case a: ArrayType => isTypeSupported(a.elementType, ARRAY_ELEMENT, fallbackReasons)
- case m: MapType =>
- isTypeSupported(m.keyType, MAP_KEY, fallbackReasons) && isTypeSupported(
- m.valueType,
- MAP_VALUE,
- fallbackReasons)
- case _ => false
+ case ByteType | ShortType if usingDataSourceExecWithIncompatTypes(SQLConf.get) =>
+ fallbackReasons += s"${CometConf.COMET_SCAN_ALLOW_INCOMPATIBLE.key} is false"
+ false
+ case _ =>
+ super.isTypeSupported(dt, name, fallbackReasons)
}
}
}
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometScanExec.scala
index 43d436122..e5fe33ffb 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometScanExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometScanExec.scala
@@ -39,13 +39,14 @@ import org.apache.spark.sql.execution.datasources._
import org.apache.spark.sql.execution.datasources.parquet.{ParquetFileFormat, ParquetOptions}
import org.apache.spark.sql.execution.datasources.v2.DataSourceRDD
import org.apache.spark.sql.execution.metric._
+import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types._
import org.apache.spark.sql.vectorized.ColumnarBatch
import org.apache.spark.util.SerializableConfiguration
import org.apache.spark.util.collection._
import org.apache.comet.{CometConf, DataTypeSupport, MetricsSupport}
-import org.apache.comet.DataTypeSupport.{ARRAY_ELEMENT, MAP_KEY, MAP_VALUE}
+import org.apache.comet.CometSparkSessionExtensions.usingDataSourceExecWithIncompatTypes
import org.apache.comet.parquet.{CometParquetFileFormat, CometParquetPartitionReaderFactory}
/**
@@ -479,28 +480,6 @@ case class CometScanExec(
object CometScanExec extends DataTypeSupport {
- override def isAdditionallySupported(
- dt: DataType,
- name: String,
- fallbackReasons: ListBuffer[String]): Boolean = {
- if (CometConf.COMET_NATIVE_SCAN_IMPL.get() == CometConf.SCAN_NATIVE_ICEBERG_COMPAT) {
- dt match {
- case s: StructType =>
- s.fields.forall(f => isTypeSupported(f.dataType, f.name, fallbackReasons))
- case a: ArrayType =>
- isTypeSupported(a.elementType, ARRAY_ELEMENT, fallbackReasons)
- case m: MapType =>
- isTypeSupported(m.keyType, MAP_KEY, fallbackReasons) && isTypeSupported(
- m.valueType,
- MAP_VALUE,
- fallbackReasons)
- case _ => false
- }
- } else {
- false
- }
- }
-
def apply(scanExec: FileSourceScanExec, session: SparkSession): CometScanExec = {
// TreeNode.mapProductIterator is protected method.
def mapProductIterator[B: ClassTag](product: Product, f: Any => B): Array[B] = {
@@ -545,4 +524,20 @@ object CometScanExec extends DataTypeSupport {
// of ParquetFileFormat.
fileFormat.getClass().equals(classOf[ParquetFileFormat])
}
+
+ override def isTypeSupported(
+ dt: DataType,
+ name: String,
+ fallbackReasons: ListBuffer[String]): Boolean = {
+ dt match {
+ case ByteType | ShortType if usingDataSourceExecWithIncompatTypes(SQLConf.get) =>
+ fallbackReasons += s"${CometConf.COMET_SCAN_ALLOW_INCOMPATIBLE.key} is false"
+ false
+ case _: StructType | _: ArrayType | _: MapType
+ if CometConf.COMET_NATIVE_SCAN_IMPL.get() != CometConf.SCAN_NATIVE_ICEBERG_COMPAT =>
+ false
+ case _ =>
+ super.isTypeSupported(dt, name, fallbackReasons)
+ }
+ }
}
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala
index a390e2b3f..ddb325fd5 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala
@@ -137,11 +137,11 @@ case class CometSparkToColumnarExec(child: SparkPlan)
}
object CometSparkToColumnarExec extends DataTypeSupport {
- override def isAdditionallySupported(
+ override def isTypeSupported(
dt: DataType,
name: String,
fallbackReasons: ListBuffer[String]): Boolean = dt match {
- case _: StructType => true
- case _ => false
+ case _: ArrayType | _: MapType => false
+ case _ => super.isTypeSupported(dt, name, fallbackReasons)
}
}
diff --git a/spark/src/test/scala/org/apache/comet/CometCastSuite.scala b/spark/src/test/scala/org/apache/comet/CometCastSuite.scala
index 12fa45287..a9f22e385 100644
--- a/spark/src/test/scala/org/apache/comet/CometCastSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometCastSuite.scala
@@ -60,7 +60,7 @@ class CometCastSuite extends CometTestBase with AdaptiveSparkPlanHelper {
private val timestampPattern = "0123456789/:T" + whitespaceChars
lazy val usingParquetExecWithIncompatTypes: Boolean =
- DataTypeSupport.usingParquetExecWithIncompatTypes(conf)
+ CometSparkSessionExtensions.usingDataSourceExecWithIncompatTypes(conf)
test("all valid cast combinations covered") {
val names = testNames
diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
index 0030ef786..394ce5f2a 100644
--- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
@@ -140,7 +140,7 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
Byte.MaxValue)
withParquetTable(path.toString, "tbl") {
val qry = "select _9 from tbl order by _11"
- if (CometSparkSessionExtensions.usingDataFusionParquetExec(conf)) {
+ if (CometSparkSessionExtensions.usingDataSourceExec(conf)) {
if (!allowIncompatible) {
checkSparkAnswer(qry)
} else {
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala
index 9ef1a43cb..b792c444b 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala
@@ -758,7 +758,7 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar
// TODO: revisit this when we have resolution of https://github.com/apache/arrow-rs/issues/7040
// and https://github.com/apache/arrow-rs/issues/7097
val fieldsToTest =
- if (CometSparkSessionExtensions.usingDataFusionParquetExec(conf)) {
+ if (CometSparkSessionExtensions.usingDataSourceExec(conf)) {
Seq(
$"_1",
$"_4",
diff --git a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala
index f485039ae..5084bd4c9 100644
--- a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala
@@ -47,9 +47,9 @@ import org.apache.spark.unsafe.types.UTF8String
import com.google.common.primitives.UnsignedLong
-import org.apache.comet.{CometConf, DataTypeSupport}
+import org.apache.comet.{CometConf, CometSparkSessionExtensions}
import org.apache.comet.CometConf.SCAN_NATIVE_ICEBERG_COMPAT
-import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, usingDataFusionParquetExec}
+import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, usingDataSourceExec}
abstract class ParquetReadSuite extends CometTestBase {
import testImplicits._
@@ -118,7 +118,7 @@ abstract class ParquetReadSuite extends CometTestBase {
val fallbackReasons = new ListBuffer[String]()
assert(CometScanExec.isTypeSupported(dt, "", fallbackReasons) == expected)
// usingDataFusionParquetExec does not support CometBatchScanExec yet
- if (!usingDataFusionParquetExec(conf)) {
+ if (!usingDataSourceExec(conf)) {
assert(CometBatchScanExec.isTypeSupported(dt, "", fallbackReasons) == expected)
}
}
@@ -169,7 +169,7 @@ abstract class ParquetReadSuite extends CometTestBase {
i.toDouble,
DateTimeUtils.toJavaDate(i))
}
- if (!DataTypeSupport.usingParquetExecWithIncompatTypes(conf)) {
+ if (!CometSparkSessionExtensions.usingDataSourceExecWithIncompatTypes(conf)) {
checkParquetScan(data)
}
checkParquetFile(data)
@@ -191,7 +191,7 @@ abstract class ParquetReadSuite extends CometTestBase {
i.toDouble,
DateTimeUtils.toJavaDate(i))
}
- if (!DataTypeSupport.usingParquetExecWithIncompatTypes(conf)) {
+ if (!CometSparkSessionExtensions.usingDataSourceExecWithIncompatTypes(conf)) {
checkParquetScan(data)
}
checkParquetFile(data)
@@ -212,7 +212,7 @@ abstract class ParquetReadSuite extends CometTestBase {
DateTimeUtils.toJavaDate(i))
}
val filter = (row: Row) => row.getBoolean(0)
- if (!DataTypeSupport.usingParquetExecWithIncompatTypes(conf)) {
+ if (!CometSparkSessionExtensions.usingDataSourceExecWithIncompatTypes(conf)) {
checkParquetScan(data, filter)
}
checkParquetFile(data, filter)
@@ -1233,7 +1233,7 @@ abstract class ParquetReadSuite extends CometTestBase {
withParquetDataFrame(data, schema = Some(readSchema)) { df =>
// TODO: validate with Spark 3.x and 'usingDataFusionParquetExec=true'
- if (enableSchemaEvolution || usingDataFusionParquetExec(conf)) {
+ if (enableSchemaEvolution || usingDataSourceExec(conf)) {
checkAnswer(df, data.map(Row.fromTuple))
} else {
assertThrows[SparkException](df.collect())
@@ -1412,7 +1412,7 @@ abstract class ParquetReadSuite extends CometTestBase {
test("row group skipping doesn't overflow when reading into larger type") {
// Spark 4.0 no longer fails for widening types SPARK-40876
// https://github.com/apache/spark/commit/3361f25dc0ff6e5233903c26ee105711b79ba967
- assume(!isSpark40Plus && !usingDataFusionParquetExec(conf))
+ assume(!isSpark40Plus && !usingDataSourceExec(conf))
withTempPath { path =>
Seq(0).toDF("a").write.parquet(path.toString)
// Reading integer 'a' as a long isn't supported. Check that an exception is raised instead
diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala
index 9257dc3d7..cb561c7ab 100644
--- a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala
@@ -432,7 +432,7 @@ abstract class CometTestBase
}
def getPrimitiveTypesParquetSchema: String = {
- if (DataTypeSupport.usingParquetExecWithIncompatTypes(conf)) {
+ if (CometSparkSessionExtensions.usingDataSourceExecWithIncompatTypes(conf)) {
// Comet complex type reader has different behavior for uint_8, uint_16 types.
// The issue stems from undefined behavior in the parquet spec and is tracked
// here: https://github.com/apache/parquet-java/issues/3142
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/ParquetDatetimeRebaseSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/ParquetDatetimeRebaseSuite.scala
index 044a99b30..1b01d2fbf 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/ParquetDatetimeRebaseSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/ParquetDatetimeRebaseSuite.scala
@@ -27,7 +27,7 @@ import org.apache.spark.sql.{CometTestBase, DataFrame, Dataset, Row}
import org.apache.spark.sql.internal.SQLConf
import org.apache.comet.{CometConf, CometSparkSessionExtensions}
-import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, usingDataFusionParquetExec}
+import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, usingDataSourceExec}
// This test checks if Comet reads ancient dates & timestamps that are before 1582, as if they are
// read according to the `LegacyBehaviorPolicy.CORRECTED` mode (i.e., no rebase) in Spark.
@@ -51,7 +51,7 @@ abstract class ParquetDatetimeRebaseSuite extends CometTestBase {
// Parquet file written by 2.4.5 should throw exception for both Spark and Comet
// For Spark 4.0+, Parquet file written by 2.4.5 should not throw exception
if ((exceptionOnRebase || sparkVersion == "2_4_5") && (!isSpark40Plus || sparkVersion != "2_4_5") && !CometSparkSessionExtensions
- .usingDataFusionParquetExec(conf)) {
+ .usingDataSourceExec(conf)) {
intercept[SparkException](df.collect())
} else {
checkSparkNoRebaseAnswer(df)
@@ -62,7 +62,7 @@ abstract class ParquetDatetimeRebaseSuite extends CometTestBase {
}
test("reading ancient timestamps before 1582") {
- assume(!usingDataFusionParquetExec(conf))
+ assume(!usingDataSourceExec(conf))
Seq(true, false).foreach { exceptionOnRebase =>
withSQLConf(
CometConf.COMET_EXCEPTION_ON_LEGACY_DATE_TIMESTAMP.key ->
@@ -76,7 +76,7 @@ abstract class ParquetDatetimeRebaseSuite extends CometTestBase {
// Parquet file written by 2.4.5 should throw exception for both Spark and Comet
// For Spark 4.0+, Parquet file written by 2.4.5 should not throw exception
if ((exceptionOnRebase || sparkVersion == "2_4_5") && (!isSpark40Plus || sparkVersion != "2_4_5") && !CometSparkSessionExtensions
- .usingDataFusionParquetExec(conf)) {
+ .usingDataSourceExec(conf)) {
intercept[SparkException](df.collect())
} else {
checkSparkNoRebaseAnswer(df)
@@ -88,7 +88,7 @@ abstract class ParquetDatetimeRebaseSuite extends CometTestBase {
}
test("reading ancient int96 timestamps before 1582") {
- assume(!usingDataFusionParquetExec(conf))
+ assume(!usingDataSourceExec(conf))
Seq(true, false).foreach { exceptionOnRebase =>
withSQLConf(
CometConf.COMET_EXCEPTION_ON_LEGACY_DATE_TIMESTAMP.key ->
@@ -102,7 +102,7 @@ abstract class ParquetDatetimeRebaseSuite extends CometTestBase {
// Parquet file written by 2.4.5 should throw exception for both Spark and Comet
// For Spark 4.0+, Parquet file written by 2.4.5 should not throw exception
if ((exceptionOnRebase || sparkVersion == "2_4_5") && (!isSpark40Plus || sparkVersion != "2_4_5") && !CometSparkSessionExtensions
- .usingDataFusionParquetExec(conf)) {
+ .usingDataSourceExec(conf)) {
intercept[SparkException](df.collect())
} else {
checkSparkNoRebaseAnswer(df)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org