Enhancement: Support table truncate and overwrite. (#34)

yuanoOo · web-flow · commit f566092b8a31 · 2025-03-31T10:47:52.000+08:00
* Enhancement: Support table truncate and overwrite.
diff --git a/spark-connector-oceanbase/spark-connector-oceanbase-3.1/src/main/scala/com/oceanbase/spark/catalog/OceanBaseTable.scala b/spark-connector-oceanbase/spark-connector-oceanbase-3.1/src/main/scala/com/oceanbase/spark/catalog/OceanBaseTable.scala
@@ -0,0 +1,74 @@
+/*
+ * Copyright 2024 OceanBase.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.oceanbase.spark.catalog
+
+import com.oceanbase.spark.config.OceanBaseConfig
+import com.oceanbase.spark.dialect.OceanBaseDialect
+import com.oceanbase.spark.reader.JDBCLimitScanBuilder
+import com.oceanbase.spark.reader.v2.OBJdbcScanBuilder
+import com.oceanbase.spark.utils.OBJdbcUtils
+import com.oceanbase.spark.writer.v2.{DirectLoadWriteBuilderV2, JDBCWriteBuilder}
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.connector.catalog._
+import org.apache.spark.sql.connector.catalog.TableCapability._
+import org.apache.spark.sql.connector.read.ScanBuilder
+import org.apache.spark.sql.connector.write.{LogicalWriteInfo, WriteBuilder}
+import org.apache.spark.sql.execution.datasources.jdbc.JDBCOptions
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+import java.util
+
+import scala.collection.JavaConverters._
+import scala.util.{Failure, Success, Try}
+
+case class OceanBaseTable(
+    ident: Identifier,
+    schema: StructType,
+    config: OceanBaseConfig,
+    dialect: OceanBaseDialect)
+  extends Table
+  with SupportsRead
+  with SupportsWrite {
+
+  override def name(): String = ident.toString
+
+  override def capabilities(): util.Set[TableCapability] = {
+    util.EnumSet.of(BATCH_READ, BATCH_WRITE, TRUNCATE)
+  }
+
+  override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = {
+    val mergedOptions = new JDBCOptions(
+      config.getProperties.asScala.toMap ++ options.asCaseSensitiveMap().asScala)
+
+    mergedOptions.parameters
+      .get(OceanBaseConfig.ENABLE_LEGACY_BATCH_READER)
+      .map(_.toBoolean) match {
+      case Some(true) => JDBCLimitScanBuilder(SparkSession.active, schema, mergedOptions)
+      case _ => OBJdbcScanBuilder(schema, config, dialect)
+    }
+  }
+
+  override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = {
+
+    if (config.getDirectLoadEnable) {
+      DirectLoadWriteBuilderV2(schema, config)
+    } else {
+      new JDBCWriteBuilder(schema, config, dialect)
+    }
+  }
+}
diff --git a/spark-connector-oceanbase/spark-connector-oceanbase-base/src/main/scala/com/oceanbase/spark/catalog/OceanBaseCatalog.scala b/spark-connector-oceanbase/spark-connector-oceanbase-base/src/main/scala/com/oceanbase/spark/catalog/OceanBaseCatalog.scala
@@ -32,6 +32,7 @@ import java.sql.SQLException
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable
+import scala.util.{Failure, Success, Try}
 
 class OceanBaseCatalog
   extends TableCatalog
@@ -181,6 +182,19 @@ class OceanBaseCatalog
     }
   }
 
+  override def purgeTable(ident: Identifier): Boolean = {
+    val config = genNewOceanBaseConfig(this.config, ident)
+    Try {
+      OBJdbcUtils.withConnection(config) {
+        conn =>
+          OBJdbcUtils.executeStatement(conn, config, dialect.getTruncateQuery(config.getDbTable))
+      }
+    } match {
+      case Success(_) => true
+      case Failure(_) => false
+    }
+  }
+
   override def namespaceExists(namespace: Array[String]): Boolean = namespace match {
     case Array(db) =>
       OBJdbcUtils.withConnection(config)(conn => dialect.schemaExists(conn, config, db))
diff --git a/spark-connector-oceanbase/spark-connector-oceanbase-base/src/main/scala/com/oceanbase/spark/catalog/OceanBaseTable.scala b/spark-connector-oceanbase/spark-connector-oceanbase-base/src/main/scala/com/oceanbase/spark/catalog/OceanBaseTable.scala
@@ -19,6 +19,7 @@ import com.oceanbase.spark.config.OceanBaseConfig
 import com.oceanbase.spark.dialect.OceanBaseDialect
 import com.oceanbase.spark.reader.JDBCLimitScanBuilder
 import com.oceanbase.spark.reader.v2.OBJdbcScanBuilder
+import com.oceanbase.spark.utils.OBJdbcUtils
 import com.oceanbase.spark.writer.v2.{DirectLoadWriteBuilderV2, JDBCWriteBuilder}
 
 import org.apache.spark.sql.SparkSession
@@ -33,6 +34,7 @@ import org.apache.spark.sql.util.CaseInsensitiveStringMap
 import java.util
 
 import scala.collection.JavaConverters._
+import scala.util.{Failure, Success, Try}
 
 case class OceanBaseTable(
     ident: Identifier,
@@ -41,7 +43,8 @@ case class OceanBaseTable(
     dialect: OceanBaseDialect)
   extends Table
   with SupportsRead
-  with SupportsWrite {
+  with SupportsWrite
+  with TruncatableTable {
 
   override def name(): String = ident.toString
 
@@ -69,4 +72,16 @@ case class OceanBaseTable(
       new JDBCWriteBuilder(schema, config, dialect)
     }
   }
+
+  override def truncateTable(): Boolean = {
+    Try {
+      OBJdbcUtils.withConnection(config) {
+        conn =>
+          OBJdbcUtils.executeStatement(conn, config, dialect.getTruncateQuery(config.getDbTable))
+      }
+    } match {
+      case Success(_) => true
+      case Failure(_) => false
+    }
+  }
 }
diff --git a/spark-connector-oceanbase/spark-connector-oceanbase-base/src/main/scala/com/oceanbase/spark/writer/v2/DirectLoadWriteBuilderV2.scala b/spark-connector-oceanbase/spark-connector-oceanbase-base/src/main/scala/com/oceanbase/spark/writer/v2/DirectLoadWriteBuilderV2.scala
@@ -15,18 +15,45 @@
  */
 package com.oceanbase.spark.writer.v2
 
+import com.oceanbase.spark.catalog.OceanBaseCatalogException
 import com.oceanbase.spark.config.OceanBaseConfig
 import com.oceanbase.spark.directload.{DirectLoader, DirectLoadUtils}
+import com.oceanbase.spark.utils.OBJdbcUtils
 
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.connector.metric.CustomMetric
-import org.apache.spark.sql.connector.write.{BatchWrite, DataWriter, DataWriterFactory, PhysicalWriteInfo, Write, WriteBuilder, WriterCommitMessage}
+import org.apache.spark.sql.connector.write.{BatchWrite, DataWriter, DataWriterFactory, PhysicalWriteInfo, SupportsOverwrite, SupportsTruncate, Write, WriteBuilder, WriterCommitMessage}
+import org.apache.spark.sql.sources.{AlwaysTrue, Filter}
 import org.apache.spark.sql.types.StructType
 
+import scala.util.{Failure, Try}
+
 /** Direct-load writing implementation based on Spark DataSource V2 API. */
 case class DirectLoadWriteBuilderV2(schema: StructType, config: OceanBaseConfig)
-  extends WriteBuilder {
+  extends WriteBuilder
+  with SupportsOverwrite
+  with SupportsTruncate {
   override def build(): Write = new DirectLoadWrite(schema, config)
+
+  override def overwrite(filters: Array[Filter]): WriteBuilder = {
+    if (filters.length == 1 && filters.head.isInstanceOf[AlwaysTrue]) {
+      Try {
+        OBJdbcUtils.withConnection(config) {
+          conn => OBJdbcUtils.executeStatement(conn, config, s"truncate table ${config.getDbTable}")
+        }
+      } match {
+        case Failure(exception) =>
+          throw OceanBaseCatalogException(
+            s"Failed to truncate table ${config.getDbTable}",
+            exception)
+        case _ =>
+      }
+    } else {
+      throw OceanBaseCatalogException(s"Currently only overwrite full data is supported.")
+    }
+
+    this
+  }
 }
 
 class DirectLoadWrite(schema: StructType, config: OceanBaseConfig) extends Write {
diff --git a/spark-connector-oceanbase/spark-connector-oceanbase-base/src/main/scala/com/oceanbase/spark/writer/v2/JDBCWriteBuilder.scala b/spark-connector-oceanbase/spark-connector-oceanbase-base/src/main/scala/com/oceanbase/spark/writer/v2/JDBCWriteBuilder.scala
@@ -15,17 +15,45 @@
  */
 package com.oceanbase.spark.writer.v2
 
+import com.oceanbase.spark.catalog.OceanBaseCatalogException
 import com.oceanbase.spark.config.OceanBaseConfig
 import com.oceanbase.spark.dialect.OceanBaseDialect
+import com.oceanbase.spark.utils.OBJdbcUtils
 
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.connector.metric.CustomMetric
-import org.apache.spark.sql.connector.write.{BatchWrite, DataWriter, DataWriterFactory, PhysicalWriteInfo, Write, WriteBuilder, WriterCommitMessage}
+import org.apache.spark.sql.connector.write.{BatchWrite, DataWriter, DataWriterFactory, PhysicalWriteInfo, SupportsOverwrite, SupportsTruncate, Write, WriteBuilder, WriterCommitMessage}
+import org.apache.spark.sql.sources.{AlwaysTrue, Filter}
 import org.apache.spark.sql.types.StructType
 
+import scala.util.{Failure, Try}
+
 class JDBCWriteBuilder(schema: StructType, config: OceanBaseConfig, dialect: OceanBaseDialect)
-  extends WriteBuilder {
+  extends WriteBuilder
+  with SupportsOverwrite
+  with SupportsTruncate {
   override def build(): Write = new JDBCWrite(schema, config, dialect)
+
+  override def overwrite(filters: Array[Filter]): WriteBuilder = {
+    if (filters.length == 1 && filters.head.isInstanceOf[AlwaysTrue]) {
+      Try {
+        OBJdbcUtils.withConnection(config) {
+          conn =>
+            OBJdbcUtils.executeStatement(conn, config, dialect.getTruncateQuery(config.getDbTable))
+        }
+      } match {
+        case Failure(exception) =>
+          throw OceanBaseCatalogException(
+            s"Failed to truncate table ${config.getDbTable}",
+            exception)
+        case _ =>
+      }
+    } else {
+      throw OceanBaseCatalogException(s"Currently only overwrite full data is supported.")
+    }
+
+    this
+  }
 }
 
 class JDBCWrite(schema: StructType, config: OceanBaseConfig, dialect: OceanBaseDialect)
diff --git a/spark-connector-oceanbase/spark-connector-oceanbase-base/src/test/scala/com/oceanbase/spark/OBCatalogMySQLITCase.scala b/spark-connector-oceanbase/spark-connector-oceanbase-base/src/test/scala/com/oceanbase/spark/OBCatalogMySQLITCase.scala
@@ -22,6 +22,8 @@ import org.apache.spark.sql.SparkSession
 import org.junit.jupiter.api.{AfterAll, AfterEach, Assertions, BeforeAll, BeforeEach, Test}
 import org.junit.jupiter.api.function.ThrowingSupplier
 
+import java.util
+
 class OBCatalogMySQLITCase extends OceanBaseMySQLTestBase {
 
   @BeforeEach
@@ -50,7 +52,7 @@ class OBCatalogMySQLITCase extends OceanBaseMySQLTestBase {
 
     session.sql("use ob;")
     insertTestData(session, "products")
-    queryAndVerifyTableData(session, "products")
+    queryAndVerifyTableData(session, "products", expected)
 
     session.stop()
   }
@@ -113,7 +115,7 @@ class OBCatalogMySQLITCase extends OceanBaseMySQLTestBase {
 
     session.sql("use ob;")
     insertTestData(session, "products_no_pri_key")
-    queryAndVerifyTableData(session, "products_no_pri_key")
+    queryAndVerifyTableData(session, "products_no_pri_key", expected)
     session.stop()
   }
 
@@ -132,7 +134,7 @@ class OBCatalogMySQLITCase extends OceanBaseMySQLTestBase {
     session.sql("use ob;")
     insertTestData(session, "products_full_pri_key")
 
-    queryAndVerifyTableData(session, "products_full_pri_key")
+    queryAndVerifyTableData(session, "products_full_pri_key", expected)
 
     session.stop()
   }
@@ -154,7 +156,11 @@ class OBCatalogMySQLITCase extends OceanBaseMySQLTestBase {
       .getOrCreate()
 
     insertTestData(session, "products")
-    queryAndVerifyTableData(session, "products")
+    queryAndVerifyTableData(session, "products", expected)
+
+    insertTestData(session, "products_no_pri_key")
+    session.sql("insert overwrite products select * from products_no_pri_key")
+    queryAndVerifyTableData(session, "products", expected)
     session.stop()
   }
 
@@ -173,7 +179,7 @@ class OBCatalogMySQLITCase extends OceanBaseMySQLTestBase {
     insertTestData(session, "products")
     // Test CTAS
     session.sql("create table test1 as select * from products")
-    queryAndVerifyTableData(session, "test1")
+    queryAndVerifyTableData(session, "test1", expected)
 
     // test bucket partition table:
     //   1. column comment test
@@ -199,7 +205,35 @@ class OBCatalogMySQLITCase extends OceanBaseMySQLTestBase {
     session.stop()
   }
 
-  private def queryAndVerifyTableData(session: SparkSession, tableName: String): Unit = {
+  @Test
+  def testTruncateAndOverWriteTable(): Unit = {
+    val session = SparkSession
+      .builder()
+      .master("local[*]")
+      .config("spark.sql.catalog.ob", OB_CATALOG_CLASS)
+      .config("spark.sql.catalog.ob.url", getJdbcUrl)
+      .config("spark.sql.catalog.ob.username", getUsername)
+      .config("spark.sql.catalog.ob.password", getPassword)
+      .config("spark.sql.catalog.ob.schema-name", getSchemaName)
+      .getOrCreate()
+
+    session.sql("use ob;")
+    insertTestData(session, "products")
+    session.sql("truncate table products")
+    val expect = new util.ArrayList[String]()
+    queryAndVerifyTableData(session, "products", expect)
+
+    insertTestData(session, "products_no_pri_key")
+    session.sql("insert overwrite products select * from products_no_pri_key")
+    queryAndVerifyTableData(session, "products", expected)
+
+    session.stop()
+  }
+
+  private def queryAndVerifyTableData(
+      session: SparkSession,
+      tableName: String,
+      expected: util.List[String]): Unit = {
     import scala.collection.JavaConverters._
     val actual = session
       .sql(s"select * from $tableName")
