Skip to content

Commit

Permalink
[Kernel] Add new tests for Timestamp_NTZ writes
Browse files Browse the repository at this point in the history
## Description

This PR builds on top of PR#4159 and adds new unit tests for timestampNtz writes.

Most of the write functionality (e.g. creation, inserts, insert with partition columns) 
are already exercised by DeltaTableWritesSuite. This PR adds a few more tests.

## How was this patch tested?

New UTs.
  • Loading branch information
anoopj committed Feb 25, 2025
1 parent 87a974e commit b5ebe6a
Showing 1 changed file with 64 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ import io.delta.kernel.expressions.Literal._
import io.delta.kernel.hook.PostCommitHook.PostCommitHookType
import io.delta.kernel.internal.{SnapshotImpl, TableConfig}
import io.delta.kernel.internal.checkpoints.CheckpointerSuite.selectSingleElement
import io.delta.kernel.internal.util.ColumnMapping
import io.delta.kernel.internal.util.{ColumnMapping, VectorUtils}
import io.delta.kernel.internal.util.SchemaUtils.casePreservingPartitionColNames
import io.delta.kernel.types._
import io.delta.kernel.types.DateType.DATE
Expand All @@ -45,6 +45,8 @@ import io.delta.kernel.types.TimestampType.TIMESTAMP
import io.delta.kernel.utils.CloseableIterable
import io.delta.kernel.utils.CloseableIterable.{emptyIterable, inMemoryIterable}

import org.apache.spark.sql.delta.DeltaTableFeatureException

/** Transaction commit in this suite IS REQUIRED TO use commitTransaction than .commit */
class DeltaTableWritesSuite extends DeltaTableWriteSuiteBase with ParquetSuiteBase {

Expand Down Expand Up @@ -1323,4 +1325,65 @@ class DeltaTableWritesSuite extends DeltaTableWriteSuiteBase with ParquetSuiteBa
assert(meta.get(ColumnMapping.COLUMN_MAPPING_PHYSICAL_NAME_KEY) == expPhyName)
}
}

test("create table TIMESTAMP_NTZ type should auto-enable the reader/writer feature") {
withTempDirAndEngine { (tablePath, engine) =>
val table = Table.forPath(engine, tablePath)
val txnBuilder = table.createTransactionBuilder(engine, testEngineInfo, CREATE_TABLE)

val testSchema = new StructType().add("tz", TimestampNTZType.TIMESTAMP_NTZ)
val txn = txnBuilder
.withSchema(engine, testSchema)
.build(engine)
val txnResult = commitTransaction(txn, engine, emptyIterable())

assert(txnResult.getVersion === 0)
val protocol = getProtocolActionFromCommit(engine, table, 0)
assert(protocol.isDefined)
assert(protocol.get.getInt(0) == 3)
assert(protocol.get.getInt(1) == 7)
assert(VectorUtils.toJavaList(protocol.get.getArray(3)).contains("timestampNtz"))
}
}

test("create table without TIMESTAMP_NTZ type should use minimum protocol") {
withTempDirAndEngine { (tablePath, engine) =>
val table = Table.forPath(engine, tablePath)
val txnBuilder = table.createTransactionBuilder(engine, testEngineInfo, CREATE_TABLE)

val txn = txnBuilder
.withSchema(engine, testSchema)
.build(engine)

assert(txn.getSchema(engine) === testSchema)
assert(txn.getPartitionColumns(engine) === Seq.empty.asJava)
val txnResult = commitTransaction(txn, engine, emptyIterable())

assert(txnResult.getVersion === 0)
val protocol = getProtocolActionFromCommit(engine, table, 0)
assert(protocol.isDefined)
assert(protocol.get.getInt(0) == 1)
assert(protocol.get.getInt(1) == 2)
}
}

test("schema evolution from Spark to add TIMESTAMP_NTZ type on a table created with kernel") {
withTempDirAndEngine { (tablePath, engine) =>
val table = Table.forPath(engine, tablePath)
val txnBuilder = table.createTransactionBuilder(engine, testEngineInfo, CREATE_TABLE)
val txn = txnBuilder
.withSchema(engine, testSchema)
.build(engine)
val txnResult = commitTransaction(txn, engine, emptyIterable())

assert(txnResult.getVersion === 0)
assertThrows[DeltaTableFeatureException] {
spark.sql("ALTER TABLE delta.`" + tablePath + "` ADD COLUMN newCol TIMESTAMP_NTZ")
}
spark.sql("ALTER TABLE delta.`" + tablePath +
"` SET TBLPROPERTIES ('delta.feature.timestampNtz' = 'supported')")
spark.sql("ALTER TABLE delta.`" + tablePath + "` ADD COLUMN newCol TIMESTAMP_NTZ")
}

}
}

0 comments on commit b5ebe6a

Please sign in to comment.