This is an automated email from the ASF dual-hosted git repository.
taklwu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hbase-connectors.git
The following commit(s) were added to refs/heads/master by this push:
new bca1b51 HBASE-30186 Port basic HBase operations to spark4
HBaseContext (#158)
bca1b51 is described below
commit bca1b512b2eb43db36dfa31695bcafc006598da9
Author: Wellington Ramos Chevreuil <[email protected]>
AuthorDate: Wed Jul 29 22:02:25 2026 +0100
HBASE-30186 Port basic HBase operations to spark4 HBaseContext (#158)
Signed-off-by: Tak Lon (Stephen) Wu <[email protected]>
---
.../hadoop/hbase/spark/HBaseConnectionCache.scala | 270 +++++++++++++++++++++
.../apache/hadoop/hbase/spark/HBaseContext.scala | 223 ++++++++++++++++-
.../hadoop/hbase/spark/HBaseRDDFunctions.scala | 173 +++++++++++++
.../hbase/spark/HBaseConnectionCacheSuite.scala | 247 +++++++++++++++++++
.../hadoop/hbase/spark/HBaseContextSuite.scala | 263 +++++++++++++++++++-
5 files changed, 1171 insertions(+), 5 deletions(-)
diff --git
a/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/HBaseConnectionCache.scala
b/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/HBaseConnectionCache.scala
new file mode 100644
index 0000000..70b4c62
--- /dev/null
+++
b/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/HBaseConnectionCache.scala
@@ -0,0 +1,270 @@
+/*
+ * 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
+
+import java.io.IOException
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.hbase.HConstants
+import org.apache.hadoop.hbase.TableName
+import org.apache.hadoop.hbase.client.Admin
+import org.apache.hadoop.hbase.client.Connection
+import org.apache.hadoop.hbase.client.ConnectionFactory
+import org.apache.hadoop.hbase.client.RegionLocator
+import org.apache.hadoop.hbase.client.Table
+import org.apache.hadoop.hbase.ipc.RpcControllerFactory
+import org.apache.hadoop.hbase.security.User
+import org.apache.hadoop.hbase.security.UserProvider
+import org.apache.yetus.audience.InterfaceAudience
+import scala.collection.mutable
+
[email protected]
+private[spark] object HBaseConnectionCache extends Logging {
+
+ val connectionMap = new mutable.HashMap[HBaseConnectionKey,
SmartConnection]()
+
+ val cacheStat = HBaseConnectionCacheStat(0, 0, 0)
+
+ // in milliseconds
+ private final val DEFAULT_TIME_OUT: Long = 10 * 60 * 1000
+ private var timeout = DEFAULT_TIME_OUT
+ private var closed: Boolean = false
+
+ var housekeepingThread = new Thread(new Runnable {
+ override def run(): Unit = {
+ while (true) {
+ try {
+ Thread.sleep(timeout)
+ } catch {
+ case e: InterruptedException =>
+ // setTimeout() and close() may interrupt the sleep and it's safe
+ // to ignore the exception
+ }
+ if (closed)
+ return
+ performHousekeeping(false)
+ }
+ }
+ })
+ housekeepingThread.setDaemon(true)
+ housekeepingThread.start()
+
+ def getStat: HBaseConnectionCacheStat = {
+ connectionMap.synchronized {
+ cacheStat.numActiveConnections = connectionMap.size
+ cacheStat.copy()
+ }
+ }
+
+ def close(): Unit = {
+ try {
+ connectionMap.synchronized {
+ if (closed)
+ return
+ closed = true
+ housekeepingThread.interrupt()
+ housekeepingThread = null
+ HBaseConnectionCache.performHousekeeping(true)
+ }
+ } catch {
+ case e: Exception => logWarning("Error in finalHouseKeeping", e)
+ }
+ }
+
+ def performHousekeeping(forceClean: Boolean): Unit = {
+ val tsNow: Long = System.currentTimeMillis()
+ connectionMap.synchronized {
+ connectionMap.foreach {
+ x =>
+ {
+ if (x._2.refCount < 0) {
+ logError(s"Bug to be fixed: negative refCount of connection
${x._2}")
+ }
+
+ if (forceClean || ((x._2.refCount <= 0) && (tsNow - x._2.timestamp
> timeout))) {
+ try {
+ x._2.connection.close()
+ } catch {
+ case e: IOException => logWarning(s"Fail to close connection
${x._2}", e)
+ }
+ connectionMap.remove(x._1)
+ }
+ }
+ }
+ }
+ }
+
+ // For testing purpose only
+ def getConnection(key: HBaseConnectionKey, conn: => Connection):
SmartConnection = {
+ connectionMap.synchronized {
+ if (closed)
+ return null
+ cacheStat.numTotalRequests += 1
+ val sc = connectionMap.getOrElseUpdate(
+ key, {
+ cacheStat.numActualConnectionsCreated += 1
+ new SmartConnection(conn)
+ })
+ sc.refCount += 1
+ sc
+ }
+ }
+
+ def getConnection(conf: Configuration): SmartConnection =
+ getConnection(new HBaseConnectionKey(conf),
ConnectionFactory.createConnection(conf))
+
+ // For testing purpose only
+ def setTimeout(to: Long): Unit = {
+ connectionMap.synchronized {
+ if (closed)
+ return
+ timeout = to
+ housekeepingThread.interrupt()
+ }
+ }
+}
+
[email protected]
+private[hbase] case class SmartConnection(
+ connection: Connection,
+ var refCount: Int = 0,
+ var timestamp: Long = 0) {
+ def getTable(tableName: TableName): Table = connection.getTable(tableName)
+ def getRegionLocator(tableName: TableName): RegionLocator =
connection.getRegionLocator(tableName)
+ def isClosed: Boolean = connection.isClosed
+ def getAdmin: Admin = connection.getAdmin
+ def close(): Unit = {
+ HBaseConnectionCache.connectionMap.synchronized {
+ refCount -= 1
+ if (refCount <= 0)
+ timestamp = System.currentTimeMillis()
+ }
+ }
+}
+
+/**
+ * Denotes a unique key to an HBase Connection instance.
+ * Please refer to 'org.apache.hadoop.hbase.client.HConnectionKey'.
+ *
+ * In essence, this class captures the properties in Configuration
+ * that may be used in the process of establishing a connection.
+ */
[email protected]
+class HBaseConnectionKey(c: Configuration) extends Logging {
+ val conf: Configuration = c
+ val CONNECTION_PROPERTIES: Array[String] = Array[String](
+ HConstants.ZOOKEEPER_QUORUM,
+ HConstants.ZOOKEEPER_ZNODE_PARENT,
+ HConstants.ZOOKEEPER_CLIENT_PORT,
+ HConstants.HBASE_CLIENT_PAUSE,
+ HConstants.HBASE_CLIENT_RETRIES_NUMBER,
+ HConstants.HBASE_RPC_TIMEOUT_KEY,
+ HConstants.HBASE_META_SCANNER_CACHING,
+ HConstants.HBASE_CLIENT_INSTANCE_ID,
+ HConstants.RPC_CODEC_CONF_KEY,
+ HConstants.USE_META_REPLICAS,
+ RpcControllerFactory.CUSTOM_CONTROLLER_CONF_KEY)
+
+ var username: String = _
+ var m_properties: mutable.HashMap[String, String] =
mutable.HashMap.empty[String, String]
+ if (conf != null) {
+ for (property <- CONNECTION_PROPERTIES) {
+ val value: String = conf.get(property)
+ if (value != null) {
+ m_properties.+=((property, value))
+ }
+ }
+ try {
+ val provider: UserProvider = UserProvider.instantiate(conf)
+ val currentUser: User = provider.getCurrent
+ if (currentUser != null) {
+ username = currentUser.getName
+ }
+ } catch {
+ case e: IOException =>
+ logWarning("Error obtaining current user, skipping username in
HBaseConnectionKey", e)
+ }
+ }
+
+ // make 'properties' immutable
+ val properties: Map[String, String] = m_properties.toMap
+
+ override def hashCode: Int = {
+ val prime: Int = 31
+ var result: Int = 1
+ if (username != null) {
+ result = username.hashCode
+ }
+ for (property <- CONNECTION_PROPERTIES) {
+ val value: Option[String] = properties.get(property)
+ if (value.isDefined) {
+ result = prime * result + value.hashCode
+ }
+ }
+ result
+ }
+
+ override def equals(obj: Any): Boolean = {
+ if (obj == null) return false
+ if (getClass ne obj.getClass) return false
+ val that: HBaseConnectionKey = obj.asInstanceOf[HBaseConnectionKey]
+ if (this.username != null && !(this.username == that.username)) {
+ return false
+ } else if (this.username == null && that.username != null) {
+ return false
+ }
+ if (this.properties == null) {
+ if (that.properties != null) {
+ return false
+ }
+ } else {
+ if (that.properties == null) {
+ return false
+ }
+ var flag: Boolean = true
+ for (property <- CONNECTION_PROPERTIES) {
+ val thisValue: Option[String] = this.properties.get(property)
+ val thatValue: Option[String] = that.properties.get(property)
+ flag = true
+ if (thisValue eq thatValue) {
+ flag = false // continue, so make flag be false
+ }
+ if (flag && (thisValue == null || !(thisValue == thatValue))) {
+ return false
+ }
+ }
+ }
+ true
+ }
+
+ override def toString: String = {
+ "HBaseConnectionKey{" + "properties=" + properties + ", username='" +
username + '\'' + '}'
+ }
+}
+
+/**
+ * To log the state of 'HBaseConnectionCache'
+ *
+ * @param numTotalRequests number of total connection requests to the cache
+ * @param numActualConnectionsCreated number of actual HBase connections the
cache ever created
+ * @param numActiveConnections number of current alive HBase connections the
cache is holding
+ */
[email protected]
+case class HBaseConnectionCacheStat(
+ var numTotalRequests: Long,
+ var numActualConnectionsCreated: Long,
+ var numActiveConnections: Long)
diff --git
a/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/HBaseContext.scala
b/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/HBaseContext.scala
index da66ed6..0a9be4b 100644
---
a/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/HBaseContext.scala
+++
b/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/HBaseContext.scala
@@ -19,7 +19,7 @@ package org.apache.hadoop.hbase.spark
import org.apache.hadoop.conf.Configuration
import org.apache.hadoop.hbase.{CellUtil, TableName}
-import org.apache.hadoop.hbase.client.{Result, Scan}
+import org.apache.hadoop.hbase.client._
import org.apache.hadoop.hbase.io.ImmutableBytesWritable
import org.apache.hadoop.hbase.mapreduce.{IdentityTableMapper,
TableInputFormat, TableMapReduceUtil}
import org.apache.hadoop.hbase.util.Bytes
@@ -34,9 +34,10 @@ import org.apache.yetus.audience.InterfaceAudience
import scala.reflect.ClassTag
/**
- * Narrow Spark 4 port of `org.apache.hadoop.hbase.spark.HBaseContext`: enough
for datasource scans
- * via `NewHBaseRDD` / TableInputFormat. Bulk load, streaming, and imperative
helpers are intentionally
- * omitted for now and will be implemented later (see HBASE-30178).
+ * HBaseContext is a facade for HBase operations like bulk put, get, delete,
and scan.
+ *
+ * HBaseContext will take the responsibilities of disseminating the
configuration information
+ * to the workers and managing the life cycle of Connections.
*/
@InterfaceAudience.Public
class HBaseContext(@transient val sc: SparkContext, @transient val config:
Configuration)
@@ -55,6 +56,119 @@ class HBaseContext(@transient val sc: SparkContext,
@transient val config: Confi
LatestHBaseContextCache.latest = this
+ /**
+ * A simple enrichment of the traditional Spark RDD foreachPartition.
+ * This function differs from the original in that it offers the
+ * developer access to a already connected Connection object
+ *
+ * Note: Do not close the Connection object. All Connection
+ * management is handled outside this method
+ *
+ * @param rdd Original RDD with data to iterate over
+ * @param f Function to be given a iterator to iterate through
+ * the RDD values and a Connection object to interact
+ * with HBase
+ */
+ def foreachPartition[T](rdd: RDD[T], f: (Iterator[T], Connection) => Unit):
Unit = {
+ rdd.foreachPartition(it => hbaseForeachPartition(broadcastedConf, it, f))
+ }
+
+ /**
+ * A simple enrichment of the traditional Spark RDD mapPartition.
+ * This function differs from the original in that it offers the
+ * developer access to a already connected Connection object
+ *
+ * Note: Do not close the Connection object. All Connection
+ * management is handled outside this method
+ *
+ * @param rdd Original RDD with data to iterate over
+ * @param mp Function to be given a iterator to iterate through
+ * the RDD values and a Connection object to interact
+ * with HBase
+ * @return Returns a new RDD generated by the user definition
+ * function just like normal mapPartition
+ */
+ def mapPartitions[T, R: ClassTag](
+ rdd: RDD[T],
+ mp: (Iterator[T], Connection) => Iterator[R]): RDD[R] = {
+ rdd.mapPartitions[R](it => hbaseMapPartition[T, R](broadcastedConf, it,
mp))
+ }
+
+ /**
+ * A simple abstraction over the HBaseContext.foreachPartition method.
+ *
+ * It allow addition support for a user to take RDD
+ * and generate puts and send them to HBase.
+ * The complexity of managing the Connection is
+ * removed from the developer
+ *
+ * @param rdd Original RDD with data to iterate over
+ * @param tableName The name of the table to put into
+ * @param f Function to convert a value in the RDD to a HBase Put
+ */
+ def bulkPut[T](rdd: RDD[T], tableName: TableName, f: (T) => Put): Unit = {
+ val tName = tableName.getName
+ rdd.foreachPartition(
+ it =>
+ hbaseForeachPartition[T](
+ broadcastedConf,
+ it,
+ (iterator, connection) => {
+ val m = connection.getBufferedMutator(TableName.valueOf(tName))
+ try {
+ iterator.foreach(t => m.mutate(f(t)))
+ m.flush()
+ } finally {
+ m.close()
+ }
+ }))
+ }
+
+ /**
+ * A simple abstraction over the HBaseContext.foreachPartition method.
+ *
+ * It allow addition support for a user to take a RDD and generate delete
+ * and send them to HBase. The complexity of managing the Connection is
+ * removed from the developer
+ *
+ * @param rdd Original RDD with data to iterate over
+ * @param tableName The name of the table to delete from
+ * @param f Function to convert a value in the RDD to a
+ * HBase Deletes
+ * @param batchSize The number of delete to batch before sending to
HBase
+ */
+ def bulkDelete[T](rdd: RDD[T], tableName: TableName, f: (T) => Delete,
batchSize: Integer): Unit =
+ {
+ bulkMutation(rdd, tableName, f, batchSize)
+ }
+
+ /**
+ * A simple abstraction over the HBaseContext.mapPartition method.
+ *
+ * It allow addition support for a user to take a RDD and generates a
+ * new RDD based on Gets and the results they bring back from HBase
+ *
+ * @param tableName The name of the table to get from
+ * @param batchSize How many gets to execute in a single batch
+ * @param rdd Original RDD with data to iterate over
+ * @param makeGet function to convert a value in the RDD to a
+ * HBase Get
+ * @param convertResult This will convert the HBase Result object to
+ * what ever the user wants to put in the resulting
+ * RDD
+ * @return new RDD with results from Get
+ */
+ def bulkGet[T, U: ClassTag](
+ tableName: TableName,
+ batchSize: Integer,
+ rdd: RDD[T],
+ makeGet: (T) => Get,
+ convertResult: (Result) => U): RDD[U] = {
+
+ val getMapPartition = new GetMapPartition(tableName, batchSize, makeGet,
convertResult)
+ rdd.mapPartitions[U](it => hbaseMapPartition[T, U](broadcastedConf, it,
getMapPartition.run))
+ }
+
/**
* Produces an RDD of type U by scanning an HBase table and applying a
transformation function
* to each (key, result) pair.
@@ -149,6 +263,107 @@ class HBaseContext(@transient val sc: SparkContext,
@transient val config: Confi
}
}
+ private def hbaseForeachPartition[T](
+ configBroadcast: Broadcast[SerializableWritable[Configuration]],
+ it: Iterator[T],
+ f: (Iterator[T], Connection) => Unit): Unit = {
+
+ val config = getConf(configBroadcast)
+ val smartConn = HBaseConnectionCache.getConnection(config)
+ if (smartConn == null) {
+ throw new IllegalStateException("HBaseConnectionCache is closed")
+ }
+ try {
+ f(it, smartConn.connection)
+ } finally {
+ smartConn.close()
+ }
+ }
+
+ private def hbaseMapPartition[K, U](
+ configBroadcast: Broadcast[SerializableWritable[Configuration]],
+ it: Iterator[K],
+ mp: (Iterator[K], Connection) => Iterator[U]): Iterator[U] = {
+
+ val config = getConf(configBroadcast)
+ val smartConn = HBaseConnectionCache.getConnection(config)
+ try {
+ mp(it, smartConn.connection)
+ } finally {
+ if (smartConn != null) smartConn.close()
+ }
+ }
+
+ private def bulkMutation[T](
+ rdd: RDD[T],
+ tableName: TableName,
+ f: (T) => Mutation,
+ batchSize: Integer): Unit = {
+
+ val tName = tableName.getName
+ rdd.foreachPartition(
+ it =>
+ hbaseForeachPartition[T](
+ broadcastedConf,
+ it,
+ (iterator, connection) => {
+ val table = connection.getTable(TableName.valueOf(tName))
+ try {
+ val mutationList = new java.util.ArrayList[Mutation]()
+ iterator.foreach { t =>
+ mutationList.add(f(t))
+ if (mutationList.size >= batchSize) {
+ table.batch(mutationList, null)
+ mutationList.clear()
+ }
+ }
+ if (mutationList.size() > 0) {
+ table.batch(mutationList, null)
+ mutationList.clear()
+ }
+ } finally {
+ table.close()
+ }
+ }))
+ }
+
+ private class GetMapPartition[T, U: ClassTag](
+ tableName: TableName,
+ batchSize: Integer,
+ makeGet: (T) => Get,
+ convertResult: (Result) => U)
+ extends Serializable {
+
+ val tName = tableName.getName
+
+ def run(iterator: Iterator[T], connection: Connection): Iterator[U] = {
+ val table = connection.getTable(TableName.valueOf(tName))
+ try {
+ val gets = new java.util.ArrayList[Get]()
+ val res = new scala.collection.mutable.ArrayBuffer[U]()
+
+ while (iterator.hasNext) {
+ gets.add(makeGet(iterator.next()))
+ if (gets.size() == batchSize) {
+ val results = table.get(gets)
+ res ++= results.iterator.map(convertResult)
+ gets.clear()
+ }
+ }
+
+ if (gets.size() > 0) {
+ val results = table.get(gets)
+ res ++= results.iterator.map(convertResult)
+ gets.clear()
+ }
+
+ res.iterator
+ } finally {
+ table.close()
+ }
+ }
+ }
+
private[hbase] def getConf(
configBroadcast: Broadcast[SerializableWritable[Configuration]]):
Configuration = {
diff --git
a/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/HBaseRDDFunctions.scala
b/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/HBaseRDDFunctions.scala
new file mode 100644
index 0000000..c0093ea
--- /dev/null
+++
b/spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/HBaseRDDFunctions.scala
@@ -0,0 +1,173 @@
+/*
+ * 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
+
+import org.apache.hadoop.hbase.TableName
+import org.apache.hadoop.hbase.client._
+import org.apache.hadoop.hbase.io.ImmutableBytesWritable
+import org.apache.spark.rdd.RDD
+import org.apache.yetus.audience.InterfaceAudience
+import scala.reflect.ClassTag
+
+/**
+ * HBaseRDDFunctions contains a set of implicit functions that can be
+ * applied to a Spark RDD so that we can easily interact with HBase
+ */
[email protected]
+object HBaseRDDFunctions {
+
+ /**
+ * These are implicit methods for a RDD that contains any type of
+ * data.
+ *
+ * @param rdd This is for rdd of any type
+ * @tparam T This is any type
+ */
+ implicit class GenericHBaseRDDFunctions[T](val rdd: RDD[T]) {
+
+ /**
+ * Implicit method that gives easy access to HBaseContext's bulk
+ * put. This will not return a new RDD. Think of it like a foreach
+ *
+ * @param hc The hbaseContext object to identify which
+ * HBase cluster connection to use
+ * @param tableName The tableName that the put will be sent to
+ * @param f The function that will turn the RDD values
+ * into HBase Put objects.
+ */
+ def hbaseBulkPut(hc: HBaseContext, tableName: TableName, f: (T) => Put):
Unit = {
+ hc.bulkPut(rdd, tableName, f)
+ }
+
+ /**
+ * Implicit method that gives easy access to HBaseContext's bulk
+ * get. This will return a new RDD. Think about it as a RDD map
+ * function. In that every RDD value will get a new value out of
+ * HBase. That new value will populate the newly generated RDD.
+ *
+ * @param hc The hbaseContext object to identify which
+ * HBase cluster connection to use
+ * @param tableName The tableName that the get will be sent to
+ * @param batchSize How many gets to execute in a single batch
+ * @param f The function that will turn the RDD values
+ * in HBase Get objects
+ * @param convertResult The function that will convert a HBase
+ * Result object into a value that will go
+ * into the resulting RDD
+ * @tparam R The type of Object that will be coming
+ * out of the resulting RDD
+ * @return A resulting RDD with type R objects
+ */
+ def hbaseBulkGet[R: ClassTag](
+ hc: HBaseContext,
+ tableName: TableName,
+ batchSize: Int,
+ f: (T) => Get,
+ convertResult: (Result) => R): RDD[R] = {
+ hc.bulkGet[T, R](tableName, batchSize, rdd, f, convertResult)
+ }
+
+ /**
+ * Implicit method that gives easy access to HBaseContext's bulk
+ * get. This will return a new RDD. Think about it as a RDD map
+ * function. In that every RDD value will get a new value out of
+ * HBase. That new value will populate the newly generated RDD.
+ *
+ * @param hc The hbaseContext object to identify which
+ * HBase cluster connection to use
+ * @param tableName The tableName that the get will be sent to
+ * @param batchSize How many gets to execute in a single batch
+ * @param f The function that will turn the RDD values
+ * in HBase Get objects
+ * @return A resulting RDD with type R objects
+ */
+ def hbaseBulkGet(
+ hc: HBaseContext,
+ tableName: TableName,
+ batchSize: Int,
+ f: (T) => Get): RDD[(ImmutableBytesWritable, Result)] = {
+ hc
+ .bulkGet[T, (ImmutableBytesWritable, Result)](
+ tableName,
+ batchSize,
+ rdd,
+ f,
+ result =>
+ if (result != null && result.getRow != null) {
+ (new ImmutableBytesWritable(result.getRow), result)
+ } else {
+ null
+ })
+ .filter(_ != null)
+ }
+
+ /**
+ * Implicit method that gives easy access to HBaseContext's bulk
+ * Delete. This will not return a new RDD.
+ *
+ * @param hc The hbaseContext object to identify which HBase
+ * cluster connection to use
+ * @param tableName The tableName that the deletes will be sent to
+ * @param f The function that will convert the RDD value into
+ * a HBase Delete Object
+ * @param batchSize The number of Deletes to be sent in a single batch
+ */
+ def hbaseBulkDelete(
+ hc: HBaseContext,
+ tableName: TableName,
+ f: (T) => Delete,
+ batchSize: Int): Unit = {
+ hc.bulkDelete(rdd, tableName, f, batchSize)
+ }
+
+ /**
+ * Implicit method that gives easy access to HBaseContext's
+ * foreachPartition method. This will act very much like a normal RDD
+ * foreach method but for the fact that you will now have a HBase
connection
+ * while iterating through the values.
+ *
+ * @param hc The hbaseContext object to identify which HBase
+ * cluster connection to use
+ * @param f This function will get an iterator for a Partition of an
+ * RDD along with a connection object to HBase
+ */
+ def hbaseForeachPartition(hc: HBaseContext, f: (Iterator[T], Connection)
=> Unit): Unit = {
+ hc.foreachPartition(rdd, f)
+ }
+
+ /**
+ * Implicit method that gives easy access to HBaseContext's
+ * mapPartitions method. This will act very much like a normal RDD
+ * map partitions method but for the fact that you will now have a
+ * HBase connection while iterating through the values
+ *
+ * @param hc The hbaseContext object to identify which HBase
+ * cluster connection to use
+ * @param f This function will get an iterator for a Partition of an
+ * RDD along with a connection object to HBase
+ * @tparam R This is the type of objects that will go into the resulting
+ * RDD
+ * @return A resulting RDD of type R
+ */
+ def hbaseMapPartitions[R: ClassTag](
+ hc: HBaseContext,
+ f: (Iterator[T], Connection) => Iterator[R]): RDD[R] = {
+ hc.mapPartitions[T, R](rdd, f)
+ }
+ }
+}
diff --git
a/spark4/hbase-spark4/src/test/scala/org/apache/hadoop/hbase/spark/HBaseConnectionCacheSuite.scala
b/spark4/hbase-spark4/src/test/scala/org/apache/hadoop/hbase/spark/HBaseConnectionCacheSuite.scala
new file mode 100644
index 0000000..f3ecffb
--- /dev/null
+++
b/spark4/hbase-spark4/src/test/scala/org/apache/hadoop/hbase/spark/HBaseConnectionCacheSuite.scala
@@ -0,0 +1,247 @@
+/*
+ * 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
+
+import java.util.concurrent.ExecutorService
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.hbase.TableName
+import org.apache.hadoop.hbase.client.{
+ Admin,
+ BufferedMutator,
+ BufferedMutatorParams,
+ Connection,
+ RegionLocator,
+ Table,
+ TableBuilder
+}
+import org.scalatest.funsuite.AnyFunSuite
+import scala.util.Random
+
+case class HBaseConnectionKeyMocker(confId: Int) extends
HBaseConnectionKey(null) {
+ override def hashCode: Int = {
+ confId
+ }
+
+ override def equals(obj: Any): Boolean = {
+ if (!obj.isInstanceOf[HBaseConnectionKeyMocker])
+ false
+ else
+ confId == obj.asInstanceOf[HBaseConnectionKeyMocker].confId
+ }
+}
+
+class ConnectionMocker extends Connection {
+ var isClosed: Boolean = false
+
+ def getRegionLocator(tableName: TableName): RegionLocator = null
+ def getConfiguration: Configuration = null
+ override def getTable(tableName: TableName): Table = null
+ override def getTable(tableName: TableName, pool: ExecutorService): Table =
null
+ def getBufferedMutator(params: BufferedMutatorParams): BufferedMutator = null
+ def getBufferedMutator(tableName: TableName): BufferedMutator = null
+ def getAdmin: Admin = null
+ def getTableBuilder(tableName: TableName, pool: ExecutorService):
TableBuilder = null
+
+ def close(): Unit = {
+ if (isClosed)
+ throw new IllegalStateException()
+ isClosed = true
+ }
+
+ def isAborted: Boolean = true
+ def abort(why: String, e: Throwable): Unit = {}
+
+ def clearRegionLocationCache(): Unit = {}
+}
+
+class HBaseConnectionCacheSuite extends AnyFunSuite with Logging {
+ /*
+ * These tests must be performed sequentially as they operate with an
+ * unique running thread and resource.
+ */
+ test("all test cases") {
+ testBasic()
+ testWithPressureWithoutClose()
+ testWithPressureWithClose()
+ }
+
+ def cleanEnv(): Unit = {
+ HBaseConnectionCache.connectionMap.synchronized {
+ HBaseConnectionCache.connectionMap.clear()
+ HBaseConnectionCache.cacheStat.numActiveConnections = 0
+ HBaseConnectionCache.cacheStat.numActualConnectionsCreated = 0
+ HBaseConnectionCache.cacheStat.numTotalRequests = 0
+ }
+ }
+
+ def testBasic(): Unit = {
+ cleanEnv()
+ HBaseConnectionCache.setTimeout(1 * 1000)
+
+ val connKeyMocker1 = new HBaseConnectionKeyMocker(1)
+ val connKeyMocker1a = new HBaseConnectionKeyMocker(1)
+ val connKeyMocker2 = new HBaseConnectionKeyMocker(2)
+
+ val c1 = HBaseConnectionCache
+ .getConnection(connKeyMocker1, new ConnectionMocker)
+
+ assert(HBaseConnectionCache.connectionMap.size === 1)
+ assert(HBaseConnectionCache.getStat.numTotalRequests === 1)
+ assert(HBaseConnectionCache.getStat.numActualConnectionsCreated === 1)
+ assert(HBaseConnectionCache.getStat.numActiveConnections === 1)
+
+ val c1a = HBaseConnectionCache
+ .getConnection(connKeyMocker1a, new ConnectionMocker)
+
+ HBaseConnectionCache.connectionMap.synchronized {
+ assert(HBaseConnectionCache.connectionMap.size === 1)
+ assert(HBaseConnectionCache.getStat.numTotalRequests === 2)
+ assert(HBaseConnectionCache.getStat.numActualConnectionsCreated === 1)
+ assert(HBaseConnectionCache.getStat.numActiveConnections === 1)
+ }
+
+ val c2 = HBaseConnectionCache
+ .getConnection(connKeyMocker2, new ConnectionMocker)
+
+ HBaseConnectionCache.connectionMap.synchronized {
+ assert(HBaseConnectionCache.connectionMap.size === 2)
+ assert(HBaseConnectionCache.getStat.numTotalRequests === 3)
+ assert(HBaseConnectionCache.getStat.numActualConnectionsCreated === 2)
+ assert(HBaseConnectionCache.getStat.numActiveConnections === 2)
+ }
+
+ c1.close()
+ HBaseConnectionCache.connectionMap.synchronized {
+ assert(HBaseConnectionCache.connectionMap.size === 2)
+ assert(HBaseConnectionCache.getStat.numActiveConnections === 2)
+ }
+
+ c1a.close()
+ HBaseConnectionCache.connectionMap.synchronized {
+ assert(HBaseConnectionCache.connectionMap.size === 2)
+ assert(HBaseConnectionCache.getStat.numActiveConnections === 2)
+ }
+
+ Thread.sleep(3 * 1000) // Leave housekeeping thread enough time
+ HBaseConnectionCache.connectionMap.synchronized {
+ assert(HBaseConnectionCache.connectionMap.size === 1)
+ assert(
+ HBaseConnectionCache.connectionMap.iterator
+ .next()
+ ._1
+ .asInstanceOf[HBaseConnectionKeyMocker]
+ .confId === 2)
+ assert(HBaseConnectionCache.getStat.numActiveConnections === 1)
+ }
+
+ c2.close()
+ }
+
+ def testWithPressureWithoutClose(): Unit = {
+ cleanEnv()
+
+ class TestThread extends Runnable {
+ override def run(): Unit = {
+ for (i <- 0 to 999) {
+ val c = HBaseConnectionCache.getConnection(
+ new HBaseConnectionKeyMocker(Random.nextInt(10)),
+ new ConnectionMocker)
+ }
+ }
+ }
+
+ HBaseConnectionCache.setTimeout(500)
+ val threads: Array[Thread] = new Array[Thread](100)
+for (i <- 0 to 99) {
+ threads.update(i, new Thread(new TestThread()))
+ threads(i).start()
+}
+try {
+ threads.foreach { x => x.join() }
+} catch {
+ case e: InterruptedException =>
+ Thread.currentThread().interrupt()
+ throw e
+}
+
+ Thread.sleep(1000)
+ HBaseConnectionCache.connectionMap.synchronized {
+ assert(HBaseConnectionCache.connectionMap.size === 10)
+ assert(HBaseConnectionCache.getStat.numTotalRequests === 100 * 1000)
+ assert(HBaseConnectionCache.getStat.numActualConnectionsCreated === 10)
+ assert(HBaseConnectionCache.getStat.numActiveConnections === 10)
+
+ var totalRc: Int = 0
+ HBaseConnectionCache.connectionMap.foreach { x => totalRc +=
x._2.refCount }
+ assert(totalRc === 100 * 1000)
+ HBaseConnectionCache.connectionMap.foreach {
+ x =>
+ {
+ x._2.refCount = 0
+ x._2.timestamp = System.currentTimeMillis() - 1000
+ }
+ }
+ }
+ Thread.sleep(1000)
+ assert(HBaseConnectionCache.connectionMap.size === 0)
+ assert(HBaseConnectionCache.getStat.numActualConnectionsCreated === 10)
+ assert(HBaseConnectionCache.getStat.numActiveConnections === 0)
+ }
+
+ def testWithPressureWithClose(): Unit = {
+ cleanEnv()
+
+ class TestThread extends Runnable {
+ override def run(): Unit = {
+ for (i <- 0 to 999) {
+ val c = HBaseConnectionCache.getConnection(
+ new HBaseConnectionKeyMocker(Random.nextInt(10)),
+ new ConnectionMocker)
+ Thread.`yield`()
+ c.close()
+ }
+ }
+ }
+
+ HBaseConnectionCache.setTimeout(3 * 1000)
+ val threads: Array[Thread] = new Array[Thread](100)
+ for (i <- threads.indices) {
+ threads.update(i, new Thread(new TestThread()))
+ threads(i).run()
+ }
+ try {
+ threads.foreach { x => x.join() }
+ } catch {
+ case e: InterruptedException => println(e.getMessage)
+ }
+
+ HBaseConnectionCache.connectionMap.synchronized {
+ assert(HBaseConnectionCache.connectionMap.size === 10)
+ assert(HBaseConnectionCache.getStat.numTotalRequests === 100 * 1000)
+ assert(HBaseConnectionCache.getStat.numActualConnectionsCreated === 10)
+ assert(HBaseConnectionCache.getStat.numActiveConnections === 10)
+ }
+
+ Thread.sleep(6 * 1000)
+ HBaseConnectionCache.connectionMap.synchronized {
+ assert(HBaseConnectionCache.connectionMap.size === 0)
+ assert(HBaseConnectionCache.getStat.numActualConnectionsCreated === 10)
+ assert(HBaseConnectionCache.getStat.numActiveConnections === 0)
+ }
+ }
+}
diff --git
a/spark4/hbase-spark4/src/test/scala/org/apache/hadoop/hbase/spark/HBaseContextSuite.scala
b/spark4/hbase-spark4/src/test/scala/org/apache/hadoop/hbase/spark/HBaseContextSuite.scala
index d6a4b98..d374d39 100644
---
a/spark4/hbase-spark4/src/test/scala/org/apache/hadoop/hbase/spark/HBaseContextSuite.scala
+++
b/spark4/hbase-spark4/src/test/scala/org/apache/hadoop/hbase/spark/HBaseContextSuite.scala
@@ -18,13 +18,14 @@
package org.apache.hadoop.hbase.spark
import org.apache.hadoop.hbase.{CellUtil, HBaseTestingUtility, TableName}
-import org.apache.hadoop.hbase.client.{ConnectionFactory, Put, Scan}
+import org.apache.hadoop.hbase.client._
import org.apache.hadoop.hbase.filter.FirstKeyOnlyFilter
import org.apache.hadoop.hbase.io.ImmutableBytesWritable
import org.apache.hadoop.hbase.util.Bytes
import org.apache.spark.{SparkConf, SparkContext}
import org.scalatest.BeforeAndAfterAll
import org.scalatest.funsuite.AnyFunSuite
+import scala.collection.mutable.ListBuffer
class HBaseContextSuite extends AnyFunSuite with BeforeAndAfterAll with
Logging {
@@ -158,4 +159,264 @@ class HBaseContextSuite extends AnyFunSuite with
BeforeAndAfterAll with Logging
assert(rows(0).getString(0) == "value1")
assert(rows(0).get(1) == null)
}
+
+ test("bulkPut to test HBase client") {
+ val rdd = sc.parallelize(
+ Array[(Array[Byte], Array[(Array[Byte], Array[Byte], Array[Byte])])](
+ (
+ Bytes.toBytes("put1"),
+ Array((Bytes.toBytes(columnFamily), Bytes.toBytes("a"),
Bytes.toBytes("foo1")))),
+ (
+ Bytes.toBytes("put2"),
+ Array((Bytes.toBytes(columnFamily), Bytes.toBytes("b"),
Bytes.toBytes("foo2")))),
+ (
+ Bytes.toBytes("put3"),
+ Array((Bytes.toBytes(columnFamily), Bytes.toBytes("c"),
Bytes.toBytes("foo3"))))))
+
+ hbaseContext.bulkPut[(Array[Byte], Array[(Array[Byte], Array[Byte],
Array[Byte])])](
+ rdd,
+ TableName.valueOf(tableName),
+ (putRecord) => {
+ val put = new Put(putRecord._1)
+ putRecord._2.foreach((putValue) => put.addColumn(putValue._1,
putValue._2, putValue._3))
+ put
+ })
+
+ val connection =
ConnectionFactory.createConnection(TEST_UTIL.getConfiguration)
+ val table = connection.getTable(TableName.valueOf(tableName))
+ try {
+ val foo1 = Bytes.toString(
+ CellUtil.cloneValue(
+ table
+ .get(new Get(Bytes.toBytes("put1")))
+ .getColumnLatestCell(Bytes.toBytes(columnFamily),
Bytes.toBytes("a"))))
+ assert(foo1 == "foo1")
+
+ val foo2 = Bytes.toString(
+ CellUtil.cloneValue(
+ table
+ .get(new Get(Bytes.toBytes("put2")))
+ .getColumnLatestCell(Bytes.toBytes(columnFamily),
Bytes.toBytes("b"))))
+ assert(foo2 == "foo2")
+
+ val foo3 = Bytes.toString(
+ CellUtil.cloneValue(
+ table
+ .get(new Get(Bytes.toBytes("put3")))
+ .getColumnLatestCell(Bytes.toBytes(columnFamily),
Bytes.toBytes("c"))))
+ assert(foo3 == "foo3")
+ } finally {
+ table.close()
+ connection.close()
+ }
+ }
+
+ test("bulkDelete to test HBase client") {
+ val connection =
ConnectionFactory.createConnection(TEST_UTIL.getConfiguration)
+ val table = connection.getTable(TableName.valueOf(tableName))
+
+ try {
+ var put = new Put(Bytes.toBytes("delete1"))
+ put.addColumn(Bytes.toBytes(columnFamily), Bytes.toBytes("a"),
Bytes.toBytes("foo1"))
+ table.put(put)
+ put = new Put(Bytes.toBytes("delete2"))
+ put.addColumn(Bytes.toBytes(columnFamily), Bytes.toBytes("a"),
Bytes.toBytes("foo2"))
+ table.put(put)
+ put = new Put(Bytes.toBytes("delete3"))
+ put.addColumn(Bytes.toBytes(columnFamily), Bytes.toBytes("a"),
Bytes.toBytes("foo3"))
+ table.put(put)
+ } finally {
+ table.close()
+ connection.close()
+ }
+
+ val rdd = sc.parallelize(Array[Array[Byte]](Bytes.toBytes("delete1"),
Bytes.toBytes("delete3")))
+
+ hbaseContext.bulkDelete[Array[Byte]](
+ rdd,
+ TableName.valueOf(tableName),
+ putRecord => new Delete(putRecord),
+ 4)
+
+ val connection2 =
ConnectionFactory.createConnection(TEST_UTIL.getConfiguration)
+ val table2 = connection2.getTable(TableName.valueOf(tableName))
+ try {
+ assert(
+ table2
+ .get(new Get(Bytes.toBytes("delete1")))
+ .getColumnLatestCell(Bytes.toBytes(columnFamily),
Bytes.toBytes("a")) == null)
+ assert(
+ table2
+ .get(new Get(Bytes.toBytes("delete3")))
+ .getColumnLatestCell(Bytes.toBytes(columnFamily),
Bytes.toBytes("a")) == null)
+ assert(
+ Bytes
+ .toString(
+ CellUtil.cloneValue(table2
+ .get(new Get(Bytes.toBytes("delete2")))
+ .getColumnLatestCell(Bytes.toBytes(columnFamily),
Bytes.toBytes("a"))))
+ .equals("foo2"))
+ } finally {
+ table2.close()
+ connection2.close()
+ }
+ }
+
+ test("bulkGet to test HBase client") {
+ val connection =
ConnectionFactory.createConnection(TEST_UTIL.getConfiguration)
+ val table = connection.getTable(TableName.valueOf(tableName))
+
+ try {
+ var put = new Put(Bytes.toBytes("get1"))
+ put.addColumn(Bytes.toBytes(columnFamily), Bytes.toBytes("a"),
Bytes.toBytes("foo1"))
+ table.put(put)
+ put = new Put(Bytes.toBytes("get2"))
+ put.addColumn(Bytes.toBytes(columnFamily), Bytes.toBytes("a"),
Bytes.toBytes("foo2"))
+ table.put(put)
+ put = new Put(Bytes.toBytes("get3"))
+ put.addColumn(Bytes.toBytes(columnFamily), Bytes.toBytes("a"),
Bytes.toBytes("foo3"))
+ table.put(put)
+ } finally {
+ table.close()
+ connection.close()
+ }
+
+ val rdd = sc.parallelize(
+ Array[Array[Byte]](
+ Bytes.toBytes("get1"),
+ Bytes.toBytes("get2"),
+ Bytes.toBytes("get3"),
+ Bytes.toBytes("get4")))
+
+ val getRdd = hbaseContext.bulkGet[Array[Byte], String](
+ TableName.valueOf(tableName),
+ 2,
+ rdd,
+ record => {
+ new Get(record)
+ },
+ (result: Result) => {
+ if (result.listCells() != null) {
+ val it = result.listCells().iterator()
+ val B = new StringBuilder
+
+ B.append(Bytes.toString(result.getRow) + ":")
+
+ while (it.hasNext) {
+ val cell = it.next()
+ val q = Bytes.toString(CellUtil.cloneQualifier(cell))
+ if (q.equals("counter")) {
+ B.append("(" + q + "," + Bytes.toLong(CellUtil.cloneValue(cell))
+ ")")
+ } else {
+ B.append("(" + q + "," +
Bytes.toString(CellUtil.cloneValue(cell)) + ")")
+ }
+ }
+ B.toString
+ } else {
+ ""
+ }
+ })
+ val getArray = getRdd.collect()
+
+ assert(getArray.length == 4)
+ assert(getArray.contains("get1:(a,foo1)"))
+ assert(getArray.contains("get2:(a,foo2)"))
+ assert(getArray.contains("get3:(a,foo3)"))
+ }
+
+ test("foreachPartition with connection") {
+ val tName = tableName
+ val cf = columnFamily
+ val rdd = sc.parallelize(
+ Array[(Array[Byte], Array[Byte])](
+ (Bytes.toBytes("fp1"), Bytes.toBytes("value_fp1")),
+ (Bytes.toBytes("fp2"), Bytes.toBytes("value_fp2")),
+ (Bytes.toBytes("fp3"), Bytes.toBytes("value_fp3"))))
+
+ hbaseContext.foreachPartition[(Array[Byte], Array[Byte])](
+ rdd,
+ (it, connection) => {
+ val m = connection.getBufferedMutator(TableName.valueOf(tName))
+ it.foreach {
+ case (rowKey, value) =>
+ val put = new Put(rowKey)
+ put.addColumn(Bytes.toBytes(cf), Bytes.toBytes("a"), value)
+ m.mutate(put)
+ }
+ m.flush()
+ m.close()
+ })
+
+ val connection =
ConnectionFactory.createConnection(TEST_UTIL.getConfiguration)
+ val table = connection.getTable(TableName.valueOf(tableName))
+ try {
+ val v1 = Bytes.toString(
+ CellUtil.cloneValue(
+ table
+ .get(new Get(Bytes.toBytes("fp1")))
+ .getColumnLatestCell(Bytes.toBytes(columnFamily),
Bytes.toBytes("a"))))
+ assert(v1 == "value_fp1")
+
+ val v2 = Bytes.toString(
+ CellUtil.cloneValue(
+ table
+ .get(new Get(Bytes.toBytes("fp2")))
+ .getColumnLatestCell(Bytes.toBytes(columnFamily),
Bytes.toBytes("a"))))
+ assert(v2 == "value_fp2")
+
+ val v3 = Bytes.toString(
+ CellUtil.cloneValue(
+ table
+ .get(new Get(Bytes.toBytes("fp3")))
+ .getColumnLatestCell(Bytes.toBytes(columnFamily),
Bytes.toBytes("a"))))
+ assert(v3 == "value_fp3")
+ } finally {
+ table.close()
+ connection.close()
+ }
+ }
+
+ test("mapPartitions with connection") {
+ val tName = tableName
+ val cf = columnFamily
+ val connection =
ConnectionFactory.createConnection(TEST_UTIL.getConfiguration)
+ val table = connection.getTable(TableName.valueOf(tName))
+ try {
+ var put = new Put(Bytes.toBytes("mp1"))
+ put.addColumn(Bytes.toBytes(cf), Bytes.toBytes("a"),
Bytes.toBytes("val_mp1"))
+ table.put(put)
+ put = new Put(Bytes.toBytes("mp2"))
+ put.addColumn(Bytes.toBytes(cf), Bytes.toBytes("a"),
Bytes.toBytes("val_mp2"))
+ table.put(put)
+ } finally {
+ table.close()
+ connection.close()
+ }
+
+ val rdd = sc.parallelize(Array[Array[Byte]](Bytes.toBytes("mp1"),
Bytes.toBytes("mp2")))
+
+ val resultRdd = hbaseContext.mapPartitions[Array[Byte], String](
+ rdd,
+ (it, conn) => {
+ val tbl = conn.getTable(TableName.valueOf(tName))
+ try {
+ val res = new ListBuffer[String]()
+ it.foreach { rowKey =>
+ val result = tbl.get(new Get(rowKey))
+ val cell = result.getColumnLatestCell(Bytes.toBytes(cf),
Bytes.toBytes("a"))
+ if (cell != null) {
+ res += Bytes.toString(result.getRow) + "=" +
Bytes.toString(CellUtil.cloneValue(cell))
+ }
+ }
+ res.iterator
+ } finally {
+ tbl.close()
+ }
+ })
+
+ val results = resultRdd.collect().sorted
+ assert(results.length == 2)
+ assert(results.contains("mp1=val_mp1"))
+ assert(results.contains("mp2=val_mp2"))
+ }
}
