This is an automated email from the ASF dual-hosted git repository.
gengliang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 6af121b79f5 [SPARK-38829][SQL] Add a configuration flag to enable
TIMESTAMP_NTZ support in Parquet data source
6af121b79f5 is described below
commit 6af121b79f5e571f10be28c1011196b56dc5d5d9
Author: Ivan Sadikov <[email protected]>
AuthorDate: Wed Apr 27 13:48:58 2022 +0800
[SPARK-38829][SQL] Add a configuration flag to enable TIMESTAMP_NTZ support
in Parquet data source
### What changes were proposed in this pull request?
This PR adds the configuration flag
`spark.sql.parquet.timestampNTZ.enabled` which allows to read and write
`TIMESTAMP_NTZ` values in Parquet and is enabled by default.
Previously the usage of `TIMESTAMP_NTZ` was hardcoded so there was no way
to restore the original behaviour of reading values as `TIMESTAMP_LTZ`. This PR
addresses the issue - when the config is disabled, all `TIMESTAMP_NTZ` values
will be read as `TIMESTAMP_LTZ` and writes would fail with an unsupported type
error (and will have to be cast to a different type).
### Why are the changes needed?
Allows users to fall back to reading Parquet types TIMESTAMP_MICROS and
TIMESTAMP_MILLIS with UTC offset disabled as `TIMESTAMP_LTZ`. Note that the
users must explicitly disable the config - by default we would read such
Parquet types as `TIMESTAMP_NTZ`.
### Does this PR introduce _any_ user-facing change?
Yes, users have the ability to disable `TIMESTAMP_NTZ` support in Parquet
data source.
### How was this patch tested?
I updated a few unit tests to check the behaviour of the config flag.
Closes #36158 from sadikovi/SPARK-38829-add-ntz-config-flag.
Authored-by: Ivan Sadikov <[email protected]>
Signed-off-by: Gengliang Wang <[email protected]>
---
.../org/apache/spark/sql/internal/SQLConf.scala | 13 ++++
.../parquet/SpecificParquetRecordReaderBase.java | 2 +
.../datasources/parquet/ParquetFileFormat.scala | 14 +++-
.../datasources/parquet/ParquetReadSupport.scala | 76 +++++++++++--------
.../datasources/parquet/ParquetRowConverter.scala | 3 +-
.../parquet/ParquetSchemaConverter.scala | 32 +++++---
.../datasources/v2/parquet/ParquetScan.scala | 3 +
.../parquet/ParquetFieldIdSchemaSuite.scala | 6 +-
.../datasources/parquet/ParquetIOSuite.scala | 85 ++++++++++++++++------
.../datasources/parquet/ParquetSchemaSuite.scala | 73 +++++++++++++++++--
10 files changed, 233 insertions(+), 74 deletions(-)
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index 49cd23851ec..0ba870d10e9 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -1071,6 +1071,17 @@ object SQLConf {
.booleanConf
.createWithDefault(false)
+ val PARQUET_TIMESTAMP_NTZ_ENABLED =
+ buildConf("spark.sql.parquet.timestampNTZ.enabled")
+ .doc(s"Enables ${TimestampTypes.TIMESTAMP_NTZ} support for Parquet reads
and writes. " +
+ s"When enabled, ${TimestampTypes.TIMESTAMP_NTZ} values are written as
Parquet timestamp " +
+ "columns with annotation isAdjustedToUTC = false and are inferred in a
similar way. " +
+ s"When disabled, such values are read as
${TimestampTypes.TIMESTAMP_LTZ} and have to be " +
+ s"converted to ${TimestampTypes.TIMESTAMP_LTZ} for writes.")
+ .version("3.4.0")
+ .booleanConf
+ .createWithDefault(true)
+
val ORC_COMPRESSION = buildConf("spark.sql.orc.compression.codec")
.doc("Sets the compression codec used when writing ORC files. If either
`compression` or " +
"`orc.compress` is specified in the table-specific options/properties,
the precedence " +
@@ -4536,6 +4547,8 @@ class SQLConf extends Serializable with Logging {
def ignoreMissingParquetFieldId: Boolean =
getConf(SQLConf.IGNORE_MISSING_PARQUET_FIELD_ID)
+ def parquetTimestampNTZEnabled: Boolean =
getConf(PARQUET_TIMESTAMP_NTZ_ENABLED)
+
def useV1Command: Boolean = getConf(SQLConf.LEGACY_USE_V1_COMMAND)
def histogramNumericPropagateInputType: Boolean =
diff --git
a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/SpecificParquetRecordReaderBase.java
b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/SpecificParquetRecordReaderBase.java
index 292a0f98af1..61aab6c5398 100644
---
a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/SpecificParquetRecordReaderBase.java
+++
b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/SpecificParquetRecordReaderBase.java
@@ -148,6 +148,7 @@ public abstract class SpecificParquetRecordReaderBase<T>
extends RecordReader<Vo
config.setBoolean(SQLConf.PARQUET_BINARY_AS_STRING().key() , false);
config.setBoolean(SQLConf.PARQUET_INT96_AS_TIMESTAMP().key(), false);
config.setBoolean(SQLConf.CASE_SENSITIVE().key(), false);
+ config.setBoolean(SQLConf.PARQUET_TIMESTAMP_NTZ_ENABLED().key(), false);
this.file = new Path(path);
long length =
this.file.getFileSystem(config).getFileStatus(this.file).getLen();
@@ -198,6 +199,7 @@ public abstract class SpecificParquetRecordReaderBase<T>
extends RecordReader<Vo
config.setBoolean(SQLConf.PARQUET_BINARY_AS_STRING().key() , false);
config.setBoolean(SQLConf.PARQUET_INT96_AS_TIMESTAMP().key(), false);
config.setBoolean(SQLConf.CASE_SENSITIVE().key(), false);
+ config.setBoolean(SQLConf.PARQUET_TIMESTAMP_NTZ_ENABLED().key(), false);
this.parquetColumn = new ParquetToSparkSchemaConverter(config)
.convertParquetColumn(requestedSchema, Option.empty());
this.sparkSchema = (StructType) parquetColumn.sparkType();
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFileFormat.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFileFormat.scala
index de0759979d5..d9b24565ccc 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFileFormat.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFileFormat.scala
@@ -124,6 +124,10 @@ class ParquetFileFormat
SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.key,
sparkSession.sessionState.conf.parquetFieldIdWriteEnabled.toString)
+ conf.set(
+ SQLConf.PARQUET_TIMESTAMP_NTZ_ENABLED.key,
+ sparkSession.sessionState.conf.parquetTimestampNTZEnabled.toString)
+
// Sets compression scheme
conf.set(ParquetOutputFormat.COMPRESSION,
parquetOptions.compressionCodecClassName)
@@ -230,6 +234,9 @@ class ParquetFileFormat
hadoopConf.setBoolean(
SQLConf.PARQUET_INT96_AS_TIMESTAMP.key,
sparkSession.sessionState.conf.isParquetINT96AsTimestamp)
+ hadoopConf.setBoolean(
+ SQLConf.PARQUET_TIMESTAMP_NTZ_ENABLED.key,
+ sparkSession.sessionState.conf.parquetTimestampNTZEnabled)
val broadcastedHadoopConf =
sparkSession.sparkContext.broadcast(new
SerializableConfiguration(hadoopConf))
@@ -417,7 +424,8 @@ object ParquetFileFormat extends Logging {
val converter = new ParquetToSparkSchemaConverter(
sparkSession.sessionState.conf.isParquetBinaryAsString,
- sparkSession.sessionState.conf.isParquetINT96AsTimestamp)
+ sparkSession.sessionState.conf.isParquetINT96AsTimestamp,
+ timestampNTZEnabled =
sparkSession.sessionState.conf.parquetTimestampNTZEnabled)
val seen = mutable.HashSet[String]()
val finalSchemas: Seq[StructType] = footers.flatMap { footer =>
@@ -513,12 +521,14 @@ object ParquetFileFormat extends Logging {
sparkSession: SparkSession): Option[StructType] = {
val assumeBinaryIsString =
sparkSession.sessionState.conf.isParquetBinaryAsString
val assumeInt96IsTimestamp =
sparkSession.sessionState.conf.isParquetINT96AsTimestamp
+ val timestampNTZEnabled =
sparkSession.sessionState.conf.parquetTimestampNTZEnabled
val reader = (files: Seq[FileStatus], conf: Configuration,
ignoreCorruptFiles: Boolean) => {
// Converter used to convert Parquet `MessageType` to Spark SQL
`StructType`
val converter = new ParquetToSparkSchemaConverter(
assumeBinaryIsString = assumeBinaryIsString,
- assumeInt96IsTimestamp = assumeInt96IsTimestamp)
+ assumeInt96IsTimestamp = assumeInt96IsTimestamp,
+ timestampNTZEnabled = timestampNTZEnabled)
readParquetFootersInParallel(conf, files, ignoreCorruptFiles)
.map(ParquetFileFormat.readSchemaFromFooter(_, converter))
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetReadSupport.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetReadSupport.scala
index 69684f9466f..1d35e9ea049 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetReadSupport.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetReadSupport.scala
@@ -130,6 +130,8 @@ object ParquetReadSupport extends Logging {
SQLConf.NESTED_SCHEMA_PRUNING_ENABLED.defaultValue.get)
val useFieldId = conf.getBoolean(SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key,
SQLConf.PARQUET_FIELD_ID_READ_ENABLED.defaultValue.get)
+ val timestampNTZEnabled =
conf.getBoolean(SQLConf.PARQUET_TIMESTAMP_NTZ_ENABLED.key,
+ SQLConf.PARQUET_TIMESTAMP_NTZ_ENABLED.defaultValue.get)
val ignoreMissingIds =
conf.getBoolean(SQLConf.IGNORE_MISSING_PARQUET_FIELD_ID.key,
SQLConf.IGNORE_MISSING_PARQUET_FIELD_ID.defaultValue.get)
@@ -150,7 +152,7 @@ object ParquetReadSupport extends Logging {
|""".stripMargin)
}
val parquetClippedSchema =
ParquetReadSupport.clipParquetSchema(parquetFileSchema,
- catalystRequestedSchema, caseSensitive, useFieldId)
+ catalystRequestedSchema, caseSensitive, useFieldId, timestampNTZEnabled)
// We pass two schema to ParquetRecordMaterializer:
// - parquetRequestedSchema: the schema of the file data we want to read
@@ -184,17 +186,6 @@ object ParquetReadSupport extends Logging {
parquetRequestedSchema
}
- /**
- * Overloaded method for backward compatibility with
- * `caseSensitive` default to `true` and `useFieldId` default to `false`
- */
- def clipParquetSchema(
- parquetSchema: MessageType,
- catalystSchema: StructType,
- caseSensitive: Boolean = true): MessageType = {
- clipParquetSchema(parquetSchema, catalystSchema, caseSensitive, useFieldId
= false)
- }
-
/**
* Tailors `parquetSchema` according to `catalystSchema` by removing column
paths don't exist
* in `catalystSchema`, and adding those only exist in `catalystSchema`.
@@ -203,9 +194,10 @@ object ParquetReadSupport extends Logging {
parquetSchema: MessageType,
catalystSchema: StructType,
caseSensitive: Boolean,
- useFieldId: Boolean): MessageType = {
+ useFieldId: Boolean,
+ timestampNTZEnabled: Boolean): MessageType = {
val clippedParquetFields = clipParquetGroupFields(
- parquetSchema.asGroupType(), catalystSchema, caseSensitive, useFieldId)
+ parquetSchema.asGroupType(), catalystSchema, caseSensitive, useFieldId,
timestampNTZEnabled)
if (clippedParquetFields.isEmpty) {
ParquetSchemaConverter.EMPTY_MESSAGE
} else {
@@ -220,21 +212,25 @@ object ParquetReadSupport extends Logging {
parquetType: Type,
catalystType: DataType,
caseSensitive: Boolean,
- useFieldId: Boolean): Type = {
+ useFieldId: Boolean,
+ timestampNTZEnabled: Boolean): Type = {
val newParquetType = catalystType match {
case t: ArrayType if !isPrimitiveCatalystType(t.elementType) =>
// Only clips array types with nested type as element type.
- clipParquetListType(parquetType.asGroupType(), t.elementType,
caseSensitive, useFieldId)
+ clipParquetListType(parquetType.asGroupType(), t.elementType,
caseSensitive, useFieldId,
+ timestampNTZEnabled)
case t: MapType
if !isPrimitiveCatalystType(t.keyType) ||
!isPrimitiveCatalystType(t.valueType) =>
// Only clips map types with nested key type or value type
clipParquetMapType(
- parquetType.asGroupType(), t.keyType, t.valueType, caseSensitive,
useFieldId)
+ parquetType.asGroupType(), t.keyType, t.valueType, caseSensitive,
useFieldId,
+ timestampNTZEnabled)
case t: StructType =>
- clipParquetGroup(parquetType.asGroupType(), t, caseSensitive,
useFieldId)
+ clipParquetGroup(
+ parquetType.asGroupType(), t, caseSensitive, useFieldId,
timestampNTZEnabled)
case _ =>
// UDTs and primitive types are not clipped. For UDTs, a clipped
version might not be able
@@ -270,7 +266,8 @@ object ParquetReadSupport extends Logging {
parquetList: GroupType,
elementType: DataType,
caseSensitive: Boolean,
- useFieldId: Boolean): Type = {
+ useFieldId: Boolean,
+ timestampNTZEnabled: Boolean): Type = {
// Precondition of this method, should only be called for lists with
nested element types.
assert(!isPrimitiveCatalystType(elementType))
@@ -278,7 +275,7 @@ object ParquetReadSupport extends Logging {
// list element type is just the group itself. Clip it.
if (parquetList.getLogicalTypeAnnotation == null &&
parquetList.isRepetition(Repetition.REPEATED)) {
- clipParquetType(parquetList, elementType, caseSensitive, useFieldId)
+ clipParquetType(parquetList, elementType, caseSensitive, useFieldId,
timestampNTZEnabled)
} else {
assert(
parquetList.getLogicalTypeAnnotation.isInstanceOf[ListLogicalTypeAnnotation],
@@ -310,14 +307,17 @@ object ParquetReadSupport extends Logging {
Types
.buildGroup(parquetList.getRepetition)
.as(LogicalTypeAnnotation.listType())
- .addField(clipParquetType(repeatedGroup, elementType, caseSensitive,
useFieldId))
+ .addField(
+ clipParquetType(
+ repeatedGroup, elementType, caseSensitive, useFieldId,
timestampNTZEnabled))
.named(parquetList.getName)
} else {
val newRepeatedGroup = Types
.repeatedGroup()
.addField(
clipParquetType(
- repeatedGroup.getType(0), elementType, caseSensitive,
useFieldId))
+ repeatedGroup.getType(0), elementType, caseSensitive, useFieldId,
+ timestampNTZEnabled))
.named(repeatedGroup.getName)
val newElementType = if (useFieldId && repeatedGroup.getId != null) {
@@ -347,7 +347,8 @@ object ParquetReadSupport extends Logging {
keyType: DataType,
valueType: DataType,
caseSensitive: Boolean,
- useFieldId: Boolean): GroupType = {
+ useFieldId: Boolean,
+ timestampNTZEnabled: Boolean): GroupType = {
// Precondition of this method, only handles maps with nested key types or
value types.
assert(!isPrimitiveCatalystType(keyType) ||
!isPrimitiveCatalystType(valueType))
@@ -359,8 +360,12 @@ object ParquetReadSupport extends Logging {
val newRepeatedGroup = Types
.repeatedGroup()
.as(repeatedGroup.getLogicalTypeAnnotation)
- .addField(clipParquetType(parquetKeyType, keyType, caseSensitive,
useFieldId))
- .addField(clipParquetType(parquetValueType, valueType, caseSensitive,
useFieldId))
+ .addField(
+ clipParquetType(
+ parquetKeyType, keyType, caseSensitive, useFieldId,
timestampNTZEnabled))
+ .addField(
+ clipParquetType(
+ parquetValueType, valueType, caseSensitive, useFieldId,
timestampNTZEnabled))
.named(repeatedGroup.getName)
if (useFieldId && repeatedGroup.getId != null) {
newRepeatedGroup.withId(repeatedGroup.getId.intValue())
@@ -388,9 +393,11 @@ object ParquetReadSupport extends Logging {
parquetRecord: GroupType,
structType: StructType,
caseSensitive: Boolean,
- useFieldId: Boolean): GroupType = {
+ useFieldId: Boolean,
+ timestampNTZEnabled: Boolean): GroupType = {
val clippedParquetFields =
- clipParquetGroupFields(parquetRecord, structType, caseSensitive,
useFieldId)
+ clipParquetGroupFields(parquetRecord, structType, caseSensitive,
useFieldId,
+ timestampNTZEnabled)
Types
.buildGroup(parquetRecord.getRepetition)
.as(parquetRecord.getLogicalTypeAnnotation)
@@ -407,9 +414,12 @@ object ParquetReadSupport extends Logging {
parquetRecord: GroupType,
structType: StructType,
caseSensitive: Boolean,
- useFieldId: Boolean): Seq[Type] = {
+ useFieldId: Boolean,
+ timestampNTZEnabled: Boolean): Seq[Type] = {
val toParquet = new SparkToParquetSchemaConverter(
- writeLegacyParquetFormat = false, useFieldId = useFieldId)
+ writeLegacyParquetFormat = false,
+ useFieldId = useFieldId,
+ timestampNTZEnabled = timestampNTZEnabled)
lazy val caseSensitiveParquetFieldMap =
parquetRecord.getFields.asScala.map(f => f.getName -> f).toMap
lazy val caseInsensitiveParquetFieldMap =
@@ -420,7 +430,7 @@ object ParquetReadSupport extends Logging {
def matchCaseSensitiveField(f: StructField): Type = {
caseSensitiveParquetFieldMap
.get(f.name)
- .map(clipParquetType(_, f.dataType, caseSensitive, useFieldId))
+ .map(clipParquetType(_, f.dataType, caseSensitive, useFieldId,
timestampNTZEnabled))
.getOrElse(toParquet.convertField(f))
}
@@ -435,7 +445,8 @@ object ParquetReadSupport extends Logging {
throw
QueryExecutionErrors.foundDuplicateFieldInCaseInsensitiveModeError(
f.name, parquetTypesString)
} else {
- clipParquetType(parquetTypes.head, f.dataType, caseSensitive,
useFieldId)
+ clipParquetType(
+ parquetTypes.head, f.dataType, caseSensitive, useFieldId,
timestampNTZEnabled)
}
}.getOrElse(toParquet.convertField(f))
}
@@ -451,7 +462,8 @@ object ParquetReadSupport extends Logging {
throw
QueryExecutionErrors.foundDuplicateFieldInFieldIdLookupModeError(
fieldId, parquetTypesString)
} else {
- clipParquetType(parquetTypes.head, f.dataType, caseSensitive,
useFieldId)
+ clipParquetType(
+ parquetTypes.head, f.dataType, caseSensitive, useFieldId,
timestampNTZEnabled)
}
}.getOrElse {
// When there is no ID match, we use a fake name to avoid a name
match by accident
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowConverter.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowConverter.scala
index a955dd6fc76..7bfc294a2d4 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowConverter.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowConverter.scala
@@ -484,7 +484,8 @@ private[parquet] class ParquetRowConverter(
// can be read as Spark's TimestampNTZ type. This is to avoid mistakes in
reading the timestamp
// values.
private def canReadAsTimestampNTZ(parquetType: Type): Boolean =
- parquetType.asPrimitiveType().getPrimitiveTypeName == INT64 &&
+ schemaConverter.isTimestampNTZEnabled() &&
+ parquetType.asPrimitiveType().getPrimitiveTypeName == INT64 &&
parquetType.getLogicalTypeAnnotation.isInstanceOf[TimestampLogicalTypeAnnotation]
&&
!parquetType.getLogicalTypeAnnotation
.asInstanceOf[TimestampLogicalTypeAnnotation].isAdjustedToUTC
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaConverter.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaConverter.scala
index 0e065f19a88..2e4041ee1c1 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaConverter.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaConverter.scala
@@ -47,23 +47,33 @@ import org.apache.spark.sql.types._
* @param assumeInt96IsTimestamp Whether unannotated INT96 fields should be
assumed to be Spark SQL
* [[TimestampType]] fields.
* @param caseSensitive Whether use case sensitive analysis when comparing
Spark catalyst read
- * schema with Parquet schema
+ * schema with Parquet schema.
+ * @param timestampNTZEnabled Whether TimestampNTZType type is enabled.
*/
class ParquetToSparkSchemaConverter(
assumeBinaryIsString: Boolean =
SQLConf.PARQUET_BINARY_AS_STRING.defaultValue.get,
assumeInt96IsTimestamp: Boolean =
SQLConf.PARQUET_INT96_AS_TIMESTAMP.defaultValue.get,
- caseSensitive: Boolean = SQLConf.CASE_SENSITIVE.defaultValue.get) {
+ caseSensitive: Boolean = SQLConf.CASE_SENSITIVE.defaultValue.get,
+ timestampNTZEnabled: Boolean =
SQLConf.PARQUET_TIMESTAMP_NTZ_ENABLED.defaultValue.get) {
def this(conf: SQLConf) = this(
assumeBinaryIsString = conf.isParquetBinaryAsString,
assumeInt96IsTimestamp = conf.isParquetINT96AsTimestamp,
- caseSensitive = conf.caseSensitiveAnalysis)
+ caseSensitive = conf.caseSensitiveAnalysis,
+ timestampNTZEnabled = conf.parquetTimestampNTZEnabled)
def this(conf: Configuration) = this(
assumeBinaryIsString =
conf.get(SQLConf.PARQUET_BINARY_AS_STRING.key).toBoolean,
assumeInt96IsTimestamp =
conf.get(SQLConf.PARQUET_INT96_AS_TIMESTAMP.key).toBoolean,
- caseSensitive = conf.get(SQLConf.CASE_SENSITIVE.key).toBoolean)
+ caseSensitive = conf.get(SQLConf.CASE_SENSITIVE.key).toBoolean,
+ timestampNTZEnabled =
conf.get(SQLConf.PARQUET_TIMESTAMP_NTZ_ENABLED.key).toBoolean)
+ /**
+ * Returns true if TIMESTAMP_NTZ type is enabled in this
ParquetToSparkSchemaConverter.
+ */
+ def isTimestampNTZEnabled(): Boolean = {
+ timestampNTZEnabled
+ }
/**
* Converts Parquet [[MessageType]] `parquetSchema` to a Spark SQL
[[StructType]].
@@ -250,7 +260,7 @@ class ParquetToSparkSchemaConverter(
}
case timestamp: TimestampLogicalTypeAnnotation
if timestamp.getUnit == TimeUnit.MICROS || timestamp.getUnit ==
TimeUnit.MILLIS =>
- if (timestamp.isAdjustedToUTC) {
+ if (timestamp.isAdjustedToUTC || !timestampNTZEnabled) {
TimestampType
} else {
TimestampNTZType
@@ -443,23 +453,27 @@ class ParquetToSparkSchemaConverter(
* @param outputTimestampType which parquet timestamp type to use when writing.
* @param useFieldId whether we should include write field id to Parquet
schema. Set this to false
* via `spark.sql.parquet.fieldId.write.enabled = false` to disable
writing field ids.
+ * @param timestampNTZEnabled whether TIMESTAMP_NTZ type support is enabled.
*/
class SparkToParquetSchemaConverter(
writeLegacyParquetFormat: Boolean =
SQLConf.PARQUET_WRITE_LEGACY_FORMAT.defaultValue.get,
outputTimestampType: SQLConf.ParquetOutputTimestampType.Value =
SQLConf.ParquetOutputTimestampType.INT96,
- useFieldId: Boolean =
SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.defaultValue.get) {
+ useFieldId: Boolean =
SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.defaultValue.get,
+ timestampNTZEnabled: Boolean =
SQLConf.PARQUET_TIMESTAMP_NTZ_ENABLED.defaultValue.get) {
def this(conf: SQLConf) = this(
writeLegacyParquetFormat = conf.writeLegacyParquetFormat,
outputTimestampType = conf.parquetOutputTimestampType,
- useFieldId = conf.parquetFieldIdWriteEnabled)
+ useFieldId = conf.parquetFieldIdWriteEnabled,
+ timestampNTZEnabled = conf.parquetTimestampNTZEnabled)
def this(conf: Configuration) = this(
writeLegacyParquetFormat =
conf.get(SQLConf.PARQUET_WRITE_LEGACY_FORMAT.key).toBoolean,
outputTimestampType = SQLConf.ParquetOutputTimestampType.withName(
conf.get(SQLConf.PARQUET_OUTPUT_TIMESTAMP_TYPE.key)),
- useFieldId =
conf.get(SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.key).toBoolean)
+ useFieldId =
conf.get(SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.key).toBoolean,
+ timestampNTZEnabled =
conf.get(SQLConf.PARQUET_TIMESTAMP_NTZ_ENABLED.key).toBoolean)
/**
* Converts a Spark SQL [[StructType]] to a Parquet [[MessageType]].
@@ -547,7 +561,7 @@ class SparkToParquetSchemaConverter(
.as(LogicalTypeAnnotation.timestampType(true,
TimeUnit.MILLIS)).named(field.name)
}
- case TimestampNTZType =>
+ case TimestampNTZType if timestampNTZEnabled =>
Types.primitive(INT64, repetition)
.as(LogicalTypeAnnotation.timestampType(false,
TimeUnit.MICROS)).named(field.name)
case BinaryType =>
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/parquet/ParquetScan.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/parquet/ParquetScan.scala
index 6b35f2406a8..ae55159c07d 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/parquet/ParquetScan.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/parquet/ParquetScan.scala
@@ -85,6 +85,9 @@ case class ParquetScan(
hadoopConf.setBoolean(
SQLConf.PARQUET_INT96_AS_TIMESTAMP.key,
sparkSession.sessionState.conf.isParquetINT96AsTimestamp)
+ hadoopConf.setBoolean(
+ SQLConf.PARQUET_TIMESTAMP_NTZ_ENABLED.key,
+ sparkSession.sessionState.conf.parquetTimestampNTZEnabled)
val broadcastedConf = sparkSession.sparkContext.broadcast(
new SerializableConfiguration(hadoopConf))
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFieldIdSchemaSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFieldIdSchemaSuite.scala
index b3babdd3a0c..400e0ce28ea 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFieldIdSchemaSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFieldIdSchemaSuite.scala
@@ -39,14 +39,16 @@ class ParquetFieldIdSchemaSuite extends ParquetSchemaTest {
catalystSchema: StructType,
expectedSchema: String,
caseSensitive: Boolean = true,
- useFieldId: Boolean = true): Unit = {
+ useFieldId: Boolean = true,
+ timestampNTZEnabled: Boolean = true): Unit = {
test(s"Clipping with field id - $testName") {
val fileSchema = MessageTypeParser.parseMessageType(parquetSchema)
val actual = ParquetReadSupport.clipParquetSchema(
fileSchema,
catalystSchema,
caseSensitive = caseSensitive,
- useFieldId = useFieldId)
+ useFieldId = useFieldId,
+ timestampNTZEnabled = timestampNTZEnabled)
// each fake name should be uniquely generated
val fakeColumnNames =
actual.getPaths.asScala.flatten.filter(_.startsWith(FAKE_COLUMN_NAME))
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
index 9e3be31423a..6560db4cb17 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
@@ -120,10 +120,12 @@ class ParquetIOSuite extends QueryTest with ParquetTest
with SharedSparkSession
}
test("SPARK-36182: TimestampNTZ") {
- val data = Seq("2021-01-01T00:00:00", "1970-07-15T01:02:03.456789")
- .map(ts => Tuple1(LocalDateTime.parse(ts)))
- withAllParquetReaders {
- checkParquetFile(data)
+ withSQLConf(SQLConf.PARQUET_TIMESTAMP_NTZ_ENABLED.key -> "true") {
+ val data = Seq("2021-01-01T00:00:00", "1970-07-15T01:02:03.456789")
+ .map(ts => Tuple1(LocalDateTime.parse(ts)))
+ withAllParquetReaders {
+ checkParquetFile(data)
+ }
}
}
@@ -155,27 +157,66 @@ class ParquetIOSuite extends QueryTest with ParquetTest
with SharedSparkSession
}
writer.close
- withAllParquetReaders {
- val df = spark.read.parquet(tablePath.toString)
- assertResult(df.schema) {
- StructType(
- StructField("timestamp_ltz_millis_depr", TimestampType, nullable
= true) ::
- StructField("timestamp_ltz_micros_depr", TimestampType, nullable
= true) ::
- StructField("timestamp_ltz_millis", TimestampType, nullable =
true) ::
- StructField("timestamp_ltz_micros", TimestampType, nullable =
true) ::
- StructField("timestamp_ntz_millis", TimestampNTZType, nullable =
true) ::
- StructField("timestamp_ntz_micros", TimestampNTZType, nullable =
true) ::
- Nil
- )
+ for (timestampNTZEnabled <- Seq(true, false)) {
+ withSQLConf(SQLConf.PARQUET_TIMESTAMP_NTZ_ENABLED.key ->
s"$timestampNTZEnabled") {
+ val timestampNTZType = if (timestampNTZEnabled) TimestampNTZType
else TimestampType
+
+ withAllParquetReaders {
+ val df = spark.read.parquet(tablePath.toString)
+ assertResult(df.schema) {
+ StructType(
+ StructField("timestamp_ltz_millis_depr", TimestampType,
nullable = true) ::
+ StructField("timestamp_ltz_micros_depr", TimestampType,
nullable = true) ::
+ StructField("timestamp_ltz_millis", TimestampType, nullable
= true) ::
+ StructField("timestamp_ltz_micros", TimestampType, nullable
= true) ::
+ StructField("timestamp_ntz_millis", timestampNTZType,
nullable = true) ::
+ StructField("timestamp_ntz_micros", timestampNTZType,
nullable = true) ::
+ Nil
+ )
+ }
+
+ val ltz_value = new java.sql.Timestamp(1000L)
+ val ntz_value = LocalDateTime.of(1970, 1, 1, 0, 0, 1)
+
+ val exp = if (timestampNTZEnabled) {
+ (0 until numRecords).map { _ =>
+ (ltz_value, ltz_value, ltz_value, ltz_value, ntz_value,
ntz_value)
+ }.toDF()
+ } else {
+ (0 until numRecords).map { _ =>
+ (ltz_value, ltz_value, ltz_value, ltz_value, ltz_value,
ltz_value)
+ }.toDF()
+ }
+
+ checkAnswer(df, exp)
+ }
}
+ }
+ }
+ }
+ }
- val exp = (0 until numRecords).map { _ =>
- val ltz_value = new java.sql.Timestamp(1000L)
- val ntz_value = LocalDateTime.of(1970, 1, 1, 0, 0, 1)
- (ltz_value, ltz_value, ltz_value, ltz_value, ntz_value, ntz_value)
- }.toDF()
+ test("Write TimestampNTZ type") {
+ // Writes should fail if timestamp_ntz support is disabled.
+ withSQLConf(SQLConf.PARQUET_TIMESTAMP_NTZ_ENABLED.key -> "false") {
+ withTempPath { dir =>
+ val data = Seq(LocalDateTime.parse("2021-01-01T00:00:00")).toDF("col")
+ val err = intercept[Exception] {
+ data.write.parquet(dir.getCanonicalPath)
+ }.getCause
+ assert(err.getMessage.contains("Unsupported data type timestamp_ntz"))
+ }
+ }
- checkAnswer(df, exp)
+ withSQLConf(SQLConf.PARQUET_TIMESTAMP_NTZ_ENABLED.key -> "true") {
+ withTempPath { dir =>
+ val data = Seq(LocalDateTime.parse("2021-01-01T00:00:00")).toDF("col")
+ data.write.parquet(dir.getCanonicalPath)
+ assertResult(spark.read.parquet(dir.getCanonicalPath).schema) {
+ StructType(
+ StructField("col", TimestampNTZType, nullable = true) ::
+ Nil
+ )
}
}
}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaSuite.scala
index 4579d06b26c..1bef75ef0d9 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaSuite.scala
@@ -64,12 +64,14 @@ abstract class ParquetSchemaTest extends ParquetTest with
SharedSparkSession {
binaryAsString: Boolean,
int96AsTimestamp: Boolean,
caseSensitive: Boolean = false,
+ timestampNTZEnabled: Boolean = true,
sparkReadSchema: Option[StructType] = None,
expectedParquetColumn: Option[ParquetColumn] = None): Unit = {
val converter = new ParquetToSparkSchemaConverter(
assumeBinaryIsString = binaryAsString,
assumeInt96IsTimestamp = int96AsTimestamp,
- caseSensitive = caseSensitive)
+ caseSensitive = caseSensitive,
+ timestampNTZEnabled = timestampNTZEnabled)
test(s"sql <= parquet: $testName") {
val actualParquetColumn = converter.convertParquetColumn(
@@ -95,10 +97,12 @@ abstract class ParquetSchemaTest extends ParquetTest with
SharedSparkSession {
parquetSchema: String,
writeLegacyParquetFormat: Boolean,
outputTimestampType: SQLConf.ParquetOutputTimestampType.Value =
- SQLConf.ParquetOutputTimestampType.INT96): Unit = {
+ SQLConf.ParquetOutputTimestampType.INT96,
+ timestampNTZEnabled: Boolean = true): Unit = {
val converter = new SparkToParquetSchemaConverter(
writeLegacyParquetFormat = writeLegacyParquetFormat,
- outputTimestampType = outputTimestampType)
+ outputTimestampType = outputTimestampType,
+ timestampNTZEnabled = timestampNTZEnabled)
test(s"sql => parquet: $testName") {
val actual = converter.convert(sqlSchema)
@@ -2237,7 +2241,62 @@ class ParquetSchemaSuite extends ParquetSchemaTest {
|}
""".stripMargin,
binaryAsString = true,
- int96AsTimestamp = int96AsTimestamp)
+ int96AsTimestamp = int96AsTimestamp,
+ timestampNTZEnabled = true)
+ }
+
+ testCatalystToParquet(
+ "TimestampNTZ Spark to Parquet conversion for complex types",
+ StructType(
+ Seq(
+ StructField("f1", TimestampNTZType),
+ StructField("f2", ArrayType(TimestampNTZType)),
+ StructField("f3", StructType(Seq(StructField("f4", TimestampNTZType))))
+ )
+ ),
+ """message spark_schema {
+ | optional int64 f1 (TIMESTAMP(MICROS,false));
+ | optional group f2 (LIST) {
+ | repeated group list {
+ | optional int64 element (TIMESTAMP(MICROS,false));
+ | }
+ | }
+ | optional group f3 {
+ | optional int64 f4 (TIMESTAMP(MICROS,false));
+ | }
+ |}
+ """.stripMargin,
+ writeLegacyParquetFormat = false,
+ timestampNTZEnabled = true)
+
+ for (timestampNTZEnabled <- Seq(true, false)) {
+ val dataType = if (timestampNTZEnabled) TimestampNTZType else TimestampType
+
+ testParquetToCatalyst(
+ "TimestampNTZ Parquet to Spark conversion for complex types, " +
+ s"timestampNTZEnabled: $timestampNTZEnabled",
+ StructType(
+ Seq(
+ StructField("f1", dataType),
+ StructField("f2", ArrayType(dataType)),
+ StructField("f3", StructType(Seq(StructField("f4", dataType))))
+ )
+ ),
+ """message spark_schema {
+ | optional int64 f1 (TIMESTAMP(MICROS,false));
+ | optional group f2 (LIST) {
+ | repeated group list {
+ | optional int64 element (TIMESTAMP(MICROS,false));
+ | }
+ | }
+ | optional group f3 {
+ | optional int64 f4 (TIMESTAMP(MICROS,false));
+ | }
+ |}
+ """.stripMargin,
+ binaryAsString = true,
+ int96AsTimestamp = false,
+ timestampNTZEnabled = timestampNTZEnabled)
}
private def testSchemaClipping(
@@ -2261,7 +2320,8 @@ class ParquetSchemaSuite extends ParquetSchemaTest {
MessageTypeParser.parseMessageType(parquetSchema),
catalystSchema,
caseSensitive,
- useFieldId = false)
+ useFieldId = false,
+ timestampNTZEnabled = true)
try {
expectedSchema.checkContains(actual)
@@ -2828,7 +2888,8 @@ class ParquetSchemaSuite extends ParquetSchemaTest {
MessageTypeParser.parseMessageType(parquetSchema),
catalystSchema,
caseSensitive = false,
- useFieldId = false)
+ useFieldId = false,
+ timestampNTZEnabled = false)
}
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
