You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@hbase.apache.org by wc...@apache.org on 2026/09/22 13:39:28 UTC
(hbase-connectors) 04/04: HBASE-30189 Implement SparkSQL support in the spark4 module using DataSource V2 API- #167
This is an automated email from the ASF dual-hosted git repository.
wchevreuil pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hbase-connectors.git
commit 63f734a76f5e284278bb35b1948abd65e58120aa
Author: Wellington Ramos Chevreuil <wc...@cloudera.com>
AuthorDate: Tue Sep 22 12:33:53 2026 +0100
HBASE-30189 Implement SparkSQL support in the spark4 module using DataSource V2 API- #167
DataSource V2 write path (#166)
Co-authored-by: Claude Code (claude-opus-4-6) <no...@anthropic.com>
Signed-off-by: Tak Lon (Stephen) Wu <ta...@apache.org>
---
.../hbase/spark/datasources/HBaseBatchWrite.scala | 108 +++++++++++++++
.../hbase/spark/datasources/HBaseDataWriter.scala | 148 +++++++++++++++++++++
.../spark/datasources/HBaseDataWriterFactory.scala | 49 +++++++
.../hbase/spark/datasources/HBaseTable.scala | 23 +++-
.../spark/datasources/HBaseWriteBuilder.scala | 40 ++++++
.../datasources/HBaseWriterCommitMessage.scala | 24 ++++
.../datasources/HBaseTableProviderSuite.scala | 106 ++++++++++++++-
7 files changed, 490 insertions(+), 8 deletions(-)
diff --git a/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/datasources/HBaseBatchWrite.scala b/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/datasources/HBaseBatchWrite.scala
new file mode 100644
index 0000000..837ea8f
--- /dev/null
+++ b/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/datasources/HBaseBatchWrite.scala
@@ -0,0 +1,108 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.hadoop.hbase.spark.datasources
+
+import org.apache.hadoop.fs.Path
+import org.apache.hadoop.hbase.{HBaseConfiguration, TableName}
+import org.apache.hadoop.hbase.client.{ColumnFamilyDescriptorBuilder, TableDescriptorBuilder}
+import org.apache.hadoop.hbase.spark.{HBaseConnectionCache, Logging}
+import org.apache.hadoop.hbase.util.Bytes
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.connector.write.{BatchWrite, DataWriterFactory, PhysicalWriteInfo,
+ WriterCommitMessage}
+import org.apache.spark.sql.types.StructType
+import org.apache.yetus.audience.InterfaceAudience
+
+/**
+ * This is a new class in the spark4 module. Implements BatchWrite for the DS V2 write path.
+ * Runs on the driver. Optionally creates the HBase table (when the "newtable" option is set),
+ * then produces an HBaseDataWriterFactory that is serialized to executors.
+ *
+ * In the spark 3 DS V1 model, table creation was in HBaseRelation.createTable() and the write
+ * was driven by InsertableRelation.insert() using saveAsHadoopDataset with TableOutputFormat.
+ *
+ * @param schema
+ * @param properties
+ */
+@InterfaceAudience.Private
+class HBaseBatchWrite(schema: StructType, properties: Map[String, String])
+ extends BatchWrite
+ with Logging {
+
+ private val catalog = HBaseTableCatalog(properties)
+
+ createTableIfNeeded()
+
+ override def createBatchWriterFactory(info: PhysicalWriteInfo): DataWriterFactory = {
+ val hadoopConf = SparkSession.active.sparkContext.hadoopConfiguration
+ val hbaseConf = HBaseConfiguration.create(hadoopConf)
+ properties.get(HBaseSparkConf.HBASE_CONFIG_LOCATION)
+ .foreach(_.split(",").foreach(r => hbaseConf.addResource(new Path(r))))
+ val wrappedConf = new SerializableConfiguration(hbaseConf)
+
+ new HBaseDataWriterFactory(schema, properties, catalog, wrappedConf)
+ }
+
+ override def useCommitCoordinator(): Boolean = false
+
+ override def commit(messages: Array[WriterCommitMessage]): Unit = {}
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {}
+
+ private def createTableIfNeeded(): Unit = {
+ val numReg = properties
+ .get(HBaseTableCatalog.newTable)
+ .map(_.toInt)
+ .getOrElse(0)
+
+ if (numReg <= 3) {
+ logInfo(s"${HBaseTableCatalog.newTable} is not defined or no larger than 3, " +
+ "skip the create table")
+ return
+ }
+
+ val startKey = Bytes.toBytes(
+ properties.getOrElse(HBaseTableCatalog.regionStart, HBaseTableCatalog.defaultRegionStart))
+ val endKey = Bytes.toBytes(
+ properties.getOrElse(HBaseTableCatalog.regionEnd, HBaseTableCatalog.defaultRegionEnd))
+
+ val hadoopConf = SparkSession.active.sparkContext.hadoopConfiguration
+ val conf = HBaseConfiguration.create(hadoopConf)
+ properties.get(HBaseSparkConf.HBASE_CONFIG_LOCATION)
+ .foreach(_.split(",").foreach(r => conf.addResource(new Path(r))))
+
+ val tName = TableName.valueOf(s"${catalog.namespace}:${catalog.name}")
+ val connection = HBaseConnectionCache.getConnection(conf)
+ val admin = connection.getAdmin
+ try {
+ if (!admin.tableExists(tName)) {
+ val tableDescBuilder = TableDescriptorBuilder.newBuilder(tName)
+ catalog.getColumnFamilies.foreach { cf =>
+ logDebug(s"add family $cf to $tName")
+ tableDescBuilder.setColumnFamily(
+ ColumnFamilyDescriptorBuilder.newBuilder(Bytes.toBytes(cf)).build())
+ }
+ val splitKeys = Bytes.split(startKey, endKey, numReg)
+ admin.createTable(tableDescBuilder.build(), splitKeys)
+ }
+ } finally {
+ admin.close()
+ connection.close()
+ }
+ }
+}
diff --git a/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/datasources/HBaseDataWriter.scala b/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/datasources/HBaseDataWriter.scala
new file mode 100644
index 0000000..595677a
--- /dev/null
+++ b/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/datasources/HBaseDataWriter.scala
@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.hadoop.hbase.spark.datasources
+
+import java.util.ArrayList
+import org.apache.hadoop.hbase.TableName
+import org.apache.hadoop.hbase.client.{Put, Table}
+import org.apache.hadoop.hbase.spark.{HBaseConnectionCache, Logging, SmartConnection}
+import org.apache.hadoop.hbase.util.Bytes
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage}
+import org.apache.spark.sql.types._
+import org.apache.yetus.audience.InterfaceAudience
+
+/**
+ * This is a new class in the spark4 module. Implements DataWriter[InternalRow] for the DS V2 write path.
+ * Each instance handles one Spark partition on an executor. Converts InternalRow to HBase Put operations
+ * and writes them via Table.put(List[Put]) with client-side batching.
+ *
+ * In the spark 3 DS V1 model, this logic was inside DefaultSource.insert() which used
+ * rdd.map(convertToPut).saveAsHadoopDataset() with the old mapred TableOutputFormat.
+ *
+ * @param schema
+ * @param properties
+ * @param catalog
+ * @param wrappedConf
+ */
+@InterfaceAudience.Private
+class HBaseDataWriter(
+ schema: StructType,
+ properties: Map[String, String],
+ catalog: HBaseTableCatalog,
+ wrappedConf: SerializableConfiguration)
+ extends DataWriter[InternalRow]
+ with Logging {
+
+ private val conf = wrappedConf.value
+ private val connection: SmartConnection = HBaseConnectionCache.getConnection(conf)
+ private val tableName = TableName.valueOf(s"${catalog.namespace}:${catalog.name}")
+ private val table: Table = connection.getTable(tableName)
+
+ private val timestamp = properties.get(HBaseSparkConf.TIMESTAMP).map(_.toLong)
+
+ private val rkFields = catalog.getRowKey
+ private val rkIdxedFields = rkFields.map { f =>
+ (schema.fieldIndex(f.colName), f)
+ }
+ private val colIdxedFields = schema.fieldNames
+ .filter(name => !rkFields.map(_.colName).contains(name))
+ .map(name => (schema.fieldIndex(name), catalog.sMap.getField(name)))
+
+ private val batchSize = properties
+ .get(HBaseSparkConf.BULKGET_SIZE)
+ .map(_.toInt)
+ .getOrElse(HBaseSparkConf.DEFAULT_BULKGET_SIZE)
+
+ private val putBuffer = new ArrayList[Put](batchSize)
+
+ override def write(record: InternalRow): Unit = {
+ val rowKeyBytes = buildRowKey(record)
+ val put = timestamp.fold(new Put(rowKeyBytes))(new Put(rowKeyBytes, _))
+
+ colIdxedFields.foreach { case (idx, field) =>
+ if (!record.isNullAt(idx)) {
+ val valueBytes = getValueBytes(record, idx, field)
+ put.addColumn(field.cfBytes, field.colBytes, valueBytes)
+ }
+ }
+
+ putBuffer.add(put)
+ if (putBuffer.size() >= batchSize) {
+ flushPuts()
+ }
+ }
+
+ override def commit(): WriterCommitMessage = {
+ flushPuts()
+ table.close()
+ connection.close()
+ HBaseWriterCommitMessage()
+ }
+
+ override def abort(): Unit = {
+ table.close()
+ connection.close()
+ }
+
+ override def close(): Unit = {}
+
+ private def flushPuts(): Unit = {
+ if (!putBuffer.isEmpty) {
+ table.put(putBuffer)
+ putBuffer.clear()
+ }
+ }
+
+ private def buildRowKey(record: InternalRow): Array[Byte] = {
+ val rowBytes = rkIdxedFields.map { case (idx, field) =>
+ require(!record.isNullAt(idx),
+ s"Row key column '${field.colName}' must not be null")
+ getValueBytes(record, idx, field)
+ }
+ val totalLen = rowBytes.foldLeft(0)(_ + _.length)
+ val result = new Array[Byte](totalLen)
+ var offset = 0
+ rowBytes.foreach { bytes =>
+ System.arraycopy(bytes, 0, result, offset, bytes.length)
+ offset += bytes.length
+ }
+ result
+ }
+
+ private val MILLIS_PER_DAY = 86400000L
+
+ private def getValueBytes(row: InternalRow, idx: Int, field: Field): Array[Byte] = {
+ field.dt match {
+ case BooleanType => Bytes.toBytes(row.getBoolean(idx))
+ case ByteType => Array(row.getByte(idx))
+ case ShortType => Bytes.toBytes(row.getShort(idx))
+ case IntegerType => Bytes.toBytes(row.getInt(idx))
+ case LongType => Bytes.toBytes(row.getLong(idx))
+ case FloatType => Bytes.toBytes(row.getFloat(idx))
+ case DoubleType => Bytes.toBytes(row.getDouble(idx))
+ case StringType => Bytes.toBytes(row.getUTF8String(idx).toString)
+ case BinaryType => row.getBinary(idx)
+ case DateType => Bytes.toBytes(row.getInt(idx).toLong * MILLIS_PER_DAY)
+ case TimestampType => Bytes.toBytes(row.getLong(idx) / 1000)
+ case dt: DecimalType =>
+ Bytes.toBytes(row.getDecimal(idx, dt.precision, dt.scale).toJavaBigDecimal)
+ case _ => throw new UnsupportedOperationException(s"unsupported data type ${field.dt}")
+ }
+ }
+}
diff --git a/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/datasources/HBaseDataWriterFactory.scala b/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/datasources/HBaseDataWriterFactory.scala
new file mode 100644
index 0000000..dd65a0c
--- /dev/null
+++ b/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/datasources/HBaseDataWriterFactory.scala
@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.hadoop.hbase.spark.datasources
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.connector.write.{DataWriter, DataWriterFactory}
+import org.apache.spark.sql.types.StructType
+import org.apache.yetus.audience.InterfaceAudience
+
+/**
+ * This is a new class in the spark4 module. Implements DataWriterFactory for the DS V2 write path.
+ * Serialized to executors. Creates one HBaseDataWriter per Spark partition.
+ *
+ * In the spark 3 DS V1 model, there was no factory. The RDD.saveAsHadoopDataset() call used
+ * TableOutputFormat to write Puts directly.
+ *
+ * @param schema
+ * @param properties
+ * @param catalog
+ * @param wrappedConf
+ */
+@InterfaceAudience.Private
+class HBaseDataWriterFactory(
+ schema: StructType,
+ properties: Map[String, String],
+ catalog: HBaseTableCatalog,
+ wrappedConf: SerializableConfiguration)
+ extends DataWriterFactory
+ with Serializable {
+
+ override def createWriter(partitionId: Int, taskId: Long): DataWriter[InternalRow] = {
+ new HBaseDataWriter(schema, properties, catalog, wrappedConf)
+ }
+}
diff --git a/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/datasources/HBaseTable.scala b/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/datasources/HBaseTable.scala
index c2fe3ff..2057c81 100644
--- a/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/datasources/HBaseTable.scala
+++ b/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/datasources/HBaseTable.scala
@@ -18,25 +18,29 @@
package org.apache.hadoop.hbase.spark.datasources
import java.util
-import org.apache.spark.sql.connector.catalog.{SupportsRead, Table, TableCapability}
+import org.apache.spark.sql.connector.catalog.{SupportsRead, SupportsWrite, Table, TableCapability}
import org.apache.spark.sql.connector.read.ScanBuilder
+import org.apache.spark.sql.connector.write.{LogicalWriteInfo, WriteBuilder}
import org.apache.spark.sql.types.StructType
import org.apache.spark.sql.util.CaseInsensitiveStringMap
import org.apache.yetus.audience.InterfaceAudience
import scala.jdk.CollectionConverters._
/**
- * This is a new class in the spark4 module. Implements Table and SupportsRead (TODO: modify once write ops are added).
- * Represents the HBase table as a Spark entity. Declares its schema (from the catalog) and capabilities (BATCH_READ).
- * When Spark wants to read, it calls newScanBuilder(). In the spark 3 DS V1 model,
- * there was no separate "table" concept, BaseRelation bundled schema and read logic together.
+ * This is a new class in the spark4 module. Implements Table, SupportsRead, and SupportsWrite.
+ * Represents the HBase table as a Spark entity. Declares its schema (from the catalog) and
+ * capabilities (BATCH_READ, BATCH_WRITE).
+ * When Spark wants to read, it calls newScanBuilder(). When Spark wants to write, it calls newWriteBuilder().
+ * In the spark 3 DS V1 model, there was no separate "table" concept, BaseRelation bundled schema and
+ * read/write logic together.
* @param tableSchema
* @param properties
*/
@InterfaceAudience.Private
class HBaseTable(tableSchema: StructType, properties: Map[String, String])
extends Table
- with SupportsRead {
+ with SupportsRead
+ with SupportsWrite {
override def name(): String = {
val catalog = HBaseTableCatalog(properties)
@@ -46,11 +50,16 @@ class HBaseTable(tableSchema: StructType, properties: Map[String, String])
override def schema(): StructType = tableSchema
override def capabilities(): util.Set[TableCapability] = {
- Set(TableCapability.BATCH_READ).asJava
+ Set(TableCapability.BATCH_READ, TableCapability.BATCH_WRITE).asJava
}
override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = {
val mergedProps = properties ++ options.asScala.toMap
new HBaseScanBuilder(tableSchema, mergedProps)
}
+
+ override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = {
+ val mergedProps = properties ++ info.options().asScala.toMap
+ new HBaseWriteBuilder(info.schema(), mergedProps)
+ }
}
diff --git a/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/datasources/HBaseWriteBuilder.scala b/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/datasources/HBaseWriteBuilder.scala
new file mode 100644
index 0000000..63a41af
--- /dev/null
+++ b/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/datasources/HBaseWriteBuilder.scala
@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.hadoop.hbase.spark.datasources
+
+import org.apache.spark.sql.connector.write.{BatchWrite, WriteBuilder}
+import org.apache.spark.sql.types.StructType
+import org.apache.yetus.audience.InterfaceAudience
+
+/**
+ * This is a new class in the spark4 module. Implements WriteBuilder for the DS V2 write path.
+ *
+ * In the spark 3 DS V1 model, there was no WriteBuilder. The write path was handled by
+ * CreatableRelationProvider.createRelation() which called InsertableRelation.insert() directly.
+ *
+ * @param schema
+ * @param properties
+ */
+@InterfaceAudience.Private
+class HBaseWriteBuilder(schema: StructType, properties: Map[String, String])
+ extends WriteBuilder {
+
+ override def buildForBatch(): BatchWrite = {
+ new HBaseBatchWrite(schema, properties)
+ }
+}
diff --git a/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/datasources/HBaseWriterCommitMessage.scala b/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/datasources/HBaseWriterCommitMessage.scala
new file mode 100644
index 0000000..d07507d
--- /dev/null
+++ b/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/datasources/HBaseWriterCommitMessage.scala
@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.hadoop.hbase.spark.datasources
+
+import org.apache.spark.sql.connector.write.WriterCommitMessage
+import org.apache.yetus.audience.InterfaceAudience
+
+@InterfaceAudience.Private
+case class HBaseWriterCommitMessage() extends WriterCommitMessage
diff --git a/spark4/hbase-spark4/src/test/scala/org/apache/hadoop/hbase/spark/datasources/HBaseTableProviderSuite.scala b/spark4/hbase-spark4/src/test/scala/org/apache/hadoop/hbase/spark/datasources/HBaseTableProviderSuite.scala
index 7268ce6..9d06596 100644
--- a/spark4/hbase-spark4/src/test/scala/org/apache/hadoop/hbase/spark/datasources/HBaseTableProviderSuite.scala
+++ b/spark4/hbase-spark4/src/test/scala/org/apache/hadoop/hbase/spark/datasources/HBaseTableProviderSuite.scala
@@ -22,7 +22,8 @@ import org.apache.hadoop.hbase.{HBaseTestingUtility, TableName}
import org.apache.hadoop.hbase.client.{ConnectionFactory, Put}
import org.apache.hadoop.hbase.spark.Logging
import org.apache.hadoop.hbase.util.Bytes
-import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.{Row, SparkSession}
+import org.apache.spark.sql.types.{StringType, StructField, StructType}
import org.scalatest.BeforeAndAfterAll
import org.scalatest.funsuite.AnyFunSuite
@@ -33,6 +34,7 @@ class HBaseTableProviderSuite extends AnyFunSuite with BeforeAndAfterAll with Lo
var configFile: File = _
val tableName = "test_provider"
+ val writeTableName = "test_write"
val columnFamily = "cf"
val numRows = 20
val numRowsWithoutName = 3
@@ -54,6 +56,8 @@ class HBaseTableProviderSuite extends AnyFunSuite with BeforeAndAfterAll with Lo
TEST_UTIL.createTable(TableName.valueOf(tableName), Bytes.toBytes(columnFamily))
logInfo(s" - created table $tableName")
+ TEST_UTIL.createTable(TableName.valueOf(writeTableName), Bytes.toBytes(columnFamily))
+ logInfo(s" - created table $writeTableName")
populateTestData()
@@ -216,4 +220,104 @@ class HBaseTableProviderSuite extends AnyFunSuite with BeforeAndAfterAll with Lo
.load()
assert(df.count() == numRows + numRowsWithoutName)
}
+
+ // --- Write path tests ---
+
+ val writeCatalog: String = s"""{
+ |"table":{"namespace":"default", "name":"$writeTableName"},
+ |"rowkey":"key",
+ |"columns":{
+ |"key":{"cf":"rowkey", "col":"key", "type":"string"},
+ |"name":{"cf":"$columnFamily", "col":"name", "type":"string"},
+ |"age":{"cf":"$columnFamily", "col":"age", "type":"string"},
+ |"salary":{"cf":"$columnFamily", "col":"salary", "type":"string"}
+ |}
+ |}""".stripMargin
+
+ private def loadWriteTable() = {
+ spark.read
+ .format("org.apache.hadoop.hbase.spark.datasources.HBaseTableProvider")
+ .option("catalog", writeCatalog)
+ .option(HBaseSparkConf.HBASE_CONFIG_LOCATION, configFile.getAbsolutePath)
+ .load()
+ }
+
+ test("write DataFrame and read back") {
+ val writeSchema = StructType(Seq(
+ StructField("key", StringType),
+ StructField("name", StringType),
+ StructField("age", StringType),
+ StructField("salary", StringType)))
+ val data = Seq(
+ Row("wrow000", "Alice", "30", "50000"),
+ Row("wrow001", "Bob", "25", "45000"),
+ Row("wrow002", "Carol", "35", "60000"))
+ val df = spark.createDataFrame(spark.sparkContext.parallelize(data), writeSchema)
+
+ df.write
+ .format("org.apache.hadoop.hbase.spark.datasources.HBaseTableProvider")
+ .option("catalog", writeCatalog)
+ .option(HBaseSparkConf.HBASE_CONFIG_LOCATION, configFile.getAbsolutePath)
+ .mode("append")
+ .save()
+
+ val result = loadWriteTable()
+ assert(result.count() == 3)
+ val rows = result.orderBy("key").collect()
+ assert(rows(0).getAs[String]("key") == "wrow000")
+ assert(rows(0).getAs[String]("name") == "Alice")
+ assert(rows(1).getAs[String]("key") == "wrow001")
+ assert(rows(1).getAs[String]("age") == "25")
+ assert(rows(2).getAs[String]("salary") == "60000")
+ }
+
+ test("write preserves null columns") {
+ val writeSchema = StructType(Seq(
+ StructField("key", StringType),
+ StructField("name", StringType),
+ StructField("age", StringType),
+ StructField("salary", StringType)))
+ val data = Seq(
+ Row("wnull000", null, "40", "70000"),
+ Row("wnull001", "Dave", null, null))
+ val df = spark.createDataFrame(spark.sparkContext.parallelize(data), writeSchema)
+
+ df.write
+ .format("org.apache.hadoop.hbase.spark.datasources.HBaseTableProvider")
+ .option("catalog", writeCatalog)
+ .option(HBaseSparkConf.HBASE_CONFIG_LOCATION, configFile.getAbsolutePath)
+ .mode("append")
+ .save()
+
+ val result = loadWriteTable().filter("key LIKE 'wnull%'").orderBy("key")
+ assert(result.count() == 2)
+ val rows = result.collect()
+ assert(rows(0).getAs[String]("key") == "wnull000")
+ assert(rows(0).getAs[String]("name") == null)
+ assert(rows(0).getAs[String]("age") == "40")
+ assert(rows(1).getAs[String]("key") == "wnull001")
+ assert(rows(1).getAs[String]("name") == "Dave")
+ assert(rows(1).getAs[String]("salary") == null)
+ }
+
+ test("write with short name 'hbase' alias") {
+ val writeSchema = StructType(Seq(
+ StructField("key", StringType),
+ StructField("name", StringType),
+ StructField("age", StringType),
+ StructField("salary", StringType)))
+ val data = Seq(Row("wshort000", "Eve", "28", "55000"))
+ val df = spark.createDataFrame(spark.sparkContext.parallelize(data), writeSchema)
+
+ df.write
+ .format("hbase")
+ .option("catalog", writeCatalog)
+ .option(HBaseSparkConf.HBASE_CONFIG_LOCATION, configFile.getAbsolutePath)
+ .mode("append")
+ .save()
+
+ val result = loadWriteTable().filter("key = 'wshort000'")
+ assert(result.count() == 1)
+ assert(result.first().getAs[String]("name") == "Eve")
+ }
}