code review fixes

c27kwan · Apr 11, 2024 · e1fafec · e1fafec
1 parent c3aac94
commit e1fafec
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 28 deletions.
diff --git a/spark/src/main/scala/io/delta/tables/DeltaColumnBuilder.scala b/spark/src/main/scala/io/delta/tables/DeltaColumnBuilder.scala
@@ -16,8 +16,9 @@
 
 package io.delta.tables
 
-import org.apache.spark.sql.delta.DeltaErrors
+import org.apache.spark.sql.delta.{DeltaErrors, IdentityColumn}
 import org.apache.spark.sql.delta.sources.DeltaSourceUtils.{GENERATION_EXPRESSION_METADATA_KEY, IDENTITY_INFO_ALLOW_EXPLICIT_INSERT, IDENTITY_INFO_START, IDENTITY_INFO_STEP}
+import org.apache.spark.sql.util.ScalaExtensions._
 
 import org.apache.spark.annotation._
 import org.apache.spark.sql.SparkSession
@@ -102,21 +103,21 @@ class DeltaColumnBuilder private[tables](
   /**
    * :: Evolving ::
    *
-   * Specify a column as an identity column that is always generated by the system, i.e., does not
-   * allow user specified values using a start of 1 and a step of 1.
+   * Specify a column as an identity column with default values that is always generated
+   * by the system (i.e. does not allow user-specified values).
    *
    * @since 3.2.0
    */
   @Evolving
   def generatedAlwaysAsIdentity(): DeltaColumnBuilder = {
-    generatedAlwaysAsIdentity(start = 1, step = 1)
+    generatedAlwaysAsIdentity(IdentityColumn.defaultStart, IdentityColumn.defaultStep)
   }
 
   /**
    * :: Evolving ::
    *
-   * Specify a column as an identity column that is always generated by the system, i.e., does not
-   * allow user specified values.
+   * Specify a column as an identity column that is always generated by the system (i.e. does not
+   * allow user-specified values).
    *
    * @param start the start value of the identity column
    * @param step the increment step of the identity column
@@ -134,20 +135,20 @@ class DeltaColumnBuilder private[tables](
   /**
    * :: Evolving ::
    *
-   * Specify a column as an identity column that allows user specified values such that the
-   * generated values start at 1 with a step of 1.
+   * Specify a column as an identity column that allows user-specified values such that the
+   * generated values use default start and step values.
    *
    * @since 3.2.0
    */
   @Evolving
   def generatedByDefaultAsIdentity(): DeltaColumnBuilder = {
-    generatedByDefaultAsIdentity(start = 1, step = 1)
+    generatedByDefaultAsIdentity(IdentityColumn.defaultStart, IdentityColumn.defaultStep)
   }
 
   /**
    * :: Evolving ::
    *
-   * Specify a column as an identity column that allows user specified values.
+   * Specify a column as an identity column that allows user-specified values.
    *
    * @param start the start value of the identity column
    * @param step the increment step of the identity column
@@ -190,7 +191,7 @@ class DeltaColumnBuilder private[tables](
       metadataBuilder.putString(GENERATION_EXPRESSION_METADATA_KEY, generationExpr.get)
     }
 
-    if (identityAllowExplicitInsert.isDefined) {
+    identityAllowExplicitInsert.ifDefined { allowExplicitInsert =>
       if (generationExpr.nonEmpty) {
         throw DeltaErrors.identityColumnWithGenerationExpression()
       }
@@ -200,7 +201,7 @@ class DeltaColumnBuilder private[tables](
       }
 
       metadataBuilder.putBoolean(
-        IDENTITY_INFO_ALLOW_EXPLICIT_INSERT, identityAllowExplicitInsert.get)
+        IDENTITY_INFO_ALLOW_EXPLICIT_INSERT, allowExplicitInsert)
       metadataBuilder.putLong(IDENTITY_INFO_START, identityStart.get)
       val step = identityStep.get
       if (step == 0L) {

diff --git a/spark/src/test/scala/org/apache/spark/sql/delta/DDLTestUtils.scala b/spark/src/test/scala/org/apache/spark/sql/delta/DDLTestUtils.scala
@@ -17,6 +17,7 @@
 package org.apache.spark.sql.delta
 
 import org.apache.spark.sql.delta.test.DeltaSQLCommandTest
+
 import org.apache.spark.sql.{QueryTest, SparkSession}
 import org.apache.spark.sql.test.SharedSparkSession
 import org.apache.spark.sql.types.{DataType, LongType, StructField}
@@ -96,21 +97,13 @@ case class IdentityColumnSpec(
   override def structField(spark: SparkSession): StructField = {
     var col = io.delta.tables.DeltaTable.columnBuilder(spark, colName)
       .dataType(dataType)
-    col = (generatedAsIdentityType, startsWith, incrementBy) match {
-      case (GeneratedAsIdentityType.GeneratedAlways, Some(start), Some(step)) =>
+    val start = startsWith.getOrElse(IdentityColumn.defaultStart.toLong)
+    val step = incrementBy.getOrElse(IdentityColumn.defaultStep.toLong)
+    col = generatedAsIdentityType match {
+      case GeneratedAsIdentityType.GeneratedAlways =>
         col.generatedAlwaysAsIdentity(start, step)
-      case (GeneratedAsIdentityType.GeneratedAlways, Some(start), None) =>
-        col.generatedAlwaysAsIdentity(start = start, step = 1L)
-      case (GeneratedAsIdentityType.GeneratedAlways, None, Some(step)) =>
-        col.generatedAlwaysAsIdentity(start = 1L, step = step)
-      case (GeneratedAsIdentityType.GeneratedAlways, None, None) => col.generatedAlwaysAsIdentity()
-      case (GeneratedAsIdentityType.GeneratedByDefault, Some(start), Some(step)) =>
+      case GeneratedAsIdentityType.GeneratedByDefault =>
         col.generatedByDefaultAsIdentity(start, step)
-      case (GeneratedAsIdentityType.GeneratedByDefault, Some(start), None) =>
-        col.generatedByDefaultAsIdentity(start = start, step = 1L)
-      case (GeneratedAsIdentityType.GeneratedByDefault, None, Some(step)) =>
-        col.generatedByDefaultAsIdentity(start = 1L, step = step)
-      case _ => col.generatedByDefaultAsIdentity()
     }
 
     comment.foreach { c =>

diff --git a/spark/src/test/scala/org/apache/spark/sql/delta/IdentityColumnSuite.scala b/spark/src/test/scala/org/apache/spark/sql/delta/IdentityColumnSuite.scala
@@ -36,7 +36,6 @@ trait IdentityColumnSuiteBase extends IdentityColumnTestUtils {
 
   import testImplicits._
   protected val tblName = "identity_test"
-
   test("Don't allow IDENTITY column in the schema if the feature is disabled") {
     withSQLConf(DeltaSQLConf.DELTA_IDENTITY_COLUMN_ENABLED.key -> "false") {
       withTable(tblName) {
@@ -69,7 +68,6 @@ trait IdentityColumnSuiteBase extends IdentityColumnTestUtils {
     builder.putLong(DeltaSourceUtils.IDENTITY_INFO_START, start)
     builder.putLong(DeltaSourceUtils.IDENTITY_INFO_STEP, step)
     colFields += StructField("id", LongType, true, builder.build())
-
     colFields += StructField("value", IntegerType)
 
     StructType(colFields.toSeq)
@@ -188,7 +186,6 @@ class IdentityColumnScalaSuite
   }
 }
 
-
 class IdentityColumnScalaIdColumnMappingSuite
   extends IdentityColumnSuiteBase
   with ScalaDDLTestUtils