Copilot commented on code in PR #158: URL: https://github.com/apache/hbase-connectors/pull/158#discussion_r3669304910
########## spark4/hbase-spark4/src/test/scala/org/apache/hadoop/hbase/spark/HBaseConnectionCacheSuite.scala: ########## @@ -0,0 +1,245 @@ +/* + * 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).run() + } + try { + threads.foreach { x => x.join() } + } catch { + case e: InterruptedException => println(e.getMessage) + } Review Comment: The InterruptedException is swallowed via println, which clears the interrupt status and can hide test failures. Re-interrupt the thread and rethrow so the test fails appropriately. This issue also appears on line 225 of the same file. ########## spark4/hbase-spark4/src/test/scala/org/apache/hadoop/hbase/spark/HBaseConnectionCacheSuite.scala: ########## @@ -0,0 +1,245 @@ +/* + * 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).run() + } Review Comment: Threads are created but executed with Thread.run(), which runs sequentially on the current thread. This means the test doesn't apply any concurrency pressure to the cache and can miss race conditions; use start() to actually run them concurrently. This issue also appears on line 221 of the same file. ########## 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) + } + } + } + } Review Comment: performHousekeeping() removes entries from connectionMap while iterating it via foreach. Mutating a mutable.HashMap during iteration can lead to skipped elements or undefined behavior; collect keys to remove first, then remove them in a second pass. ########## spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/HBaseContext.scala: ########## @@ -149,6 +260,99 @@ 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) + try { + f(it, smartConn.connection) + } finally { + if (smartConn != null) smartConn.close() + } Review Comment: smartConn can be null when HBaseConnectionCache is closed, but smartConn.connection is dereferenced unconditionally, which will throw a NullPointerException. Fail fast with a clear exception before dereferencing, and always close the SmartConnection in finally. This issue also appears on line 282 of the same file. ########## spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/HBaseContext.scala: ########## @@ -55,6 +56,116 @@ 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)) + iterator.foreach(T => m.mutate(f(T))) + m.flush() + m.close() Review Comment: BufferedMutator should be closed in a finally block so it isn't leaked if f() or mutate() throws. Also avoid using an uppercase identifier (T) for a value to reduce confusion with the type parameter. ########## spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/HBaseContext.scala: ########## @@ -149,6 +260,99 @@ 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) + try { + f(it, smartConn.connection) + } finally { + if (smartConn != null) 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)) + 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() + } + table.close() Review Comment: The HBase Table instance should be closed in a finally block so it isn't leaked if f() or batch() throws. This also renames the per-record value from T to t to avoid confusion with the type parameter. ########## spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/HBaseContext.scala: ########## @@ -149,6 +260,99 @@ 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) + try { + f(it, smartConn.connection) + } finally { + if (smartConn != null) 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)) + 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() + } + 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)) + + val gets = new java.util.ArrayList[Get]() + var res = List[U]() + + while (iterator.hasNext) { + gets.add(makeGet(iterator.next())) + + if (gets.size() == batchSize) { + val results = table.get(gets) + res = res ++ results.map(convertResult) + gets.clear() + } + } + if (gets.size() > 0) { + val results = table.get(gets) + res = res ++ results.map(convertResult) + gets.clear() + } + table.close() + res.iterator + } Review Comment: GetMapPartition.run builds results with repeated List concatenation (res = res ++ ...), which is O(n^2) and materializes all results in memory before returning. Use a mutable buffer (or stream results) and ensure the Table is closed in a finally block so it's not leaked on exceptions. ########## spark4/hbase-spark4/src/main/scala/org/apache/hadoop/hbase/spark/HBaseRDDFunctions.scala: ########## @@ -0,0 +1,171 @@ +/* + * 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 + }) + } Review Comment: This hbaseBulkGet overload returns null elements when a Result is missing/null. Nulls inside an RDD are easy to trip over (NPEs in downstream transformations); filter them out (or return an Option) before returning the RDD. ########## spark4/hbase-spark4/src/test/scala/org/apache/hadoop/hbase/spark/HBaseContextSuite.scala: ########## @@ -158,4 +159,261 @@ 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)) + 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)) + } + } + tbl.close() + res.iterator Review Comment: The table handle opened inside mapPartitions isn't protected by a finally; if an exception is thrown while iterating/reading, the Table can leak and make the test flaky. Wrap table usage in try/finally and close it reliably. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
