hbase-spark/src/main/scala/org/apache/hadoop/hbase/spark/BulkLoadPartitioner.scala

   1 /*
   2  * Licensed to the Apache Software Foundation (ASF) under one or more
   3  * contributor license agreements.  See the NOTICE file distributed with
   4  * this work for additional information regarding copyright ownership.
   5  * The ASF licenses this file to You under the Apache License, Version 2.0
   6  * (the "License"); you may not use this file except in compliance with
   7  * the License.  You may obtain a copy of the License at
   8  *
   9  *    http://www.apache.org/licenses/LICENSE-2.0
  10  *
  11  * Unless required by applicable law or agreed to in writing, software
  12  * distributed under the License is distributed on an "AS IS" BASIS,
  13  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14  * See the License for the specific language governing permissions and
  15  * limitations under the License.
  16  */
  17
  18 package org.apache.hadoop.hbase.spark
  19
  20 import java.util
  21 import java.util.Comparator
  22
  23 import org.apache.hadoop.hbase.util.Bytes
  24 import org.apache.spark.Partitioner
  25
  26 /**
  27  * A Partitioner implementation that will separate records to different
  28  * HBase Regions based on region splits
  29  *
  30  * @param startKeys   The start keys for the given table
  31  */
  32 class BulkLoadPartitioner(startKeys:Array[Array[Byte]])
  33   extends Partitioner {
  34
  35   override def numPartitions: Int = startKeys.length
  36
  37   override def getPartition(key: Any): Int = {
  38
  39     val comparator: Comparator[Array[Byte]] = new Comparator[Array[Byte]] {
  40       override def compare(o1: Array[Byte], o2: Array[Byte]): Int = {
  41         Bytes.compareTo(o1, o2)
  42       }
  43     }
  44
  45     val rowKey:Array[Byte] =
  46       key match {
  47         case qualifier: KeyFamilyQualifier =>
  48           qualifier.rowKey
  49         case wrapper: ByteArrayWrapper =>
  50           wrapper.value
  51         case _ =>
  52           key.asInstanceOf[Array[Byte]]
  53       }
  54     val partition = util.Arrays.binarySearch(startKeys, rowKey, comparator)
  55     if (partition < 0) partition * -1 + -2
  56     else partition
  57   }
  58 }