自定义ArangoRDD
This commit is contained in:
wanglihui
@@ -63,6 +63,18 @@
|
||||
<version>6.6.3</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>com.arangodb</groupId>
|
||||
<artifactId>velocypack-module-jdk8</artifactId>
|
||||
<version>1.1.0</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>com.arangodb</groupId>
|
||||
<artifactId>velocypack-module-scala_2.11</artifactId>
|
||||
<version>1.2.0</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.scala-lang</groupId>
|
||||
<artifactId>scala-library</artifactId>
|
||||
@@ -75,6 +87,11 @@
|
||||
<version>3.2.0</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.scala-lang.modules</groupId>
|
||||
<artifactId>scala-xml_2.11</artifactId>
|
||||
<version>1.0.4</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.scala-tools</groupId>
|
||||
|
||||
@@ -21,7 +21,7 @@ arangoDB.port=8529
|
||||
arangoDB.user=upsert
|
||||
arangoDB.password=ceiec2018
|
||||
#arangoDB.DB.name=insert_iplearn_index
|
||||
arangoDB.DB.name=ip-learning-test-0
|
||||
arangoDB.DB.name=iplearn_media_domain
|
||||
arangoDB.ttl=3600
|
||||
|
||||
thread.pool.number=5
|
||||
|
||||
@@ -0,0 +1,131 @@
|
||||
/*
|
||||
* DISCLAIMER
|
||||
*
|
||||
* Copyright 2016 ArangoDB GmbH, Cologne, Germany
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*
|
||||
* Copyright holder is ArangoDB GmbH, Cologne, Germany
|
||||
*
|
||||
* author Mark - mark at arangodb.com
|
||||
*/
|
||||
|
||||
package cn.ac.iie.spark
|
||||
|
||||
import cn.ac.iie.spark.rdd.{ArangoRdd, ReadOptions, WriteOptions}
|
||||
import cn.ac.iie.spark.vpack.VPackUtils
|
||||
import org.apache.spark.SparkContext
|
||||
import org.apache.spark.rdd.RDD
|
||||
import org.apache.spark.sql.{DataFrame, Dataset, Row}
|
||||
|
||||
import scala.collection.JavaConverters.seqAsJavaListConverter
|
||||
import scala.reflect.ClassTag
|
||||
|
||||
object ArangoSpark {
|
||||
|
||||
/**
|
||||
* Save data from rdd into ArangoDB
|
||||
*
|
||||
* @param rdd the rdd with the data to save
|
||||
* @param collection the collection to save in
|
||||
*/
|
||||
def save[T](rdd: RDD[T], collection: String): Unit =
|
||||
save(rdd, collection, WriteOptions())
|
||||
|
||||
/**
|
||||
* Save data from rdd into ArangoDB
|
||||
*
|
||||
* @param rdd the rdd with the data to save
|
||||
* @param collection the collection to save in
|
||||
* @param options additional write options
|
||||
*/
|
||||
def save[T](rdd: RDD[T], collection: String, options: WriteOptions): Unit =
|
||||
saveRDD(rdd, collection, options, (x: Iterator[T]) => x)
|
||||
|
||||
/**
|
||||
* Save data from dataset into ArangoDB
|
||||
*
|
||||
* @param dataset the dataset with data to save
|
||||
* @param collection the collection to save in
|
||||
*/
|
||||
def save[T](dataset: Dataset[T], collection: String): Unit =
|
||||
saveRDD(dataset.rdd, collection, WriteOptions(), (x: Iterator[T]) => x)
|
||||
|
||||
/**
|
||||
* Save data from dataset into ArangoDB
|
||||
*
|
||||
* @param dataset the dataset with data to save
|
||||
* @param collection the collection to save in
|
||||
* @param options additional write options
|
||||
*/
|
||||
def save[T](dataset: Dataset[T], collection: String, options: WriteOptions): Unit =
|
||||
saveRDD(dataset.rdd, collection, options, (x: Iterator[T]) => x)
|
||||
|
||||
/**
|
||||
* Save data from dataframe into ArangoDB
|
||||
*
|
||||
* @param dataframe the dataframe with data to save
|
||||
* @param collection the collection to save in
|
||||
* @param options additional write options
|
||||
*/
|
||||
def saveDF(dataframe: DataFrame, collection: String): Unit =
|
||||
saveRDD[Row](dataframe.rdd, collection, WriteOptions(), (x: Iterator[Row]) => x.map { y => VPackUtils.rowToVPack(y) })
|
||||
|
||||
/**
|
||||
* Save data from dataframe into ArangoDB
|
||||
*
|
||||
* @param dataframe the dataframe with data to save
|
||||
* @param collection the collection to save in
|
||||
* @param options additional write options
|
||||
*/
|
||||
def saveDF(dataframe: DataFrame, collection: String, options: WriteOptions): Unit =
|
||||
saveRDD[Row](dataframe.rdd, collection, options, (x: Iterator[Row]) => x.map { y => VPackUtils.rowToVPack(y) })
|
||||
|
||||
private def saveRDD[T](rdd: RDD[T], collection: String, options: WriteOptions, map: Iterator[T] => Iterator[Any]): Unit = {
|
||||
val writeOptions = createWriteOptions(options, rdd.sparkContext.getConf)
|
||||
rdd.foreachPartition { p =>
|
||||
if (p.nonEmpty) {
|
||||
val arangoDB = createArangoBuilder(writeOptions).build()
|
||||
val col = arangoDB.db(writeOptions.database).collection(collection)
|
||||
val docs = map(p).toList.asJava
|
||||
writeOptions.method match {
|
||||
case WriteOptions.INSERT => col.insertDocuments(docs)
|
||||
case WriteOptions.UPDATE => col.updateDocuments(docs)
|
||||
case WriteOptions.REPLACE => col.replaceDocuments(docs)
|
||||
}
|
||||
|
||||
arangoDB.shutdown()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Load data from ArangoDB into rdd
|
||||
*
|
||||
* @param sparkContext the sparkContext containing the ArangoDB configuration
|
||||
* @param collection the collection to load data from
|
||||
*/
|
||||
def load[T: ClassTag](sparkContext: SparkContext, collection: String): ArangoRdd[T] =
|
||||
load(sparkContext, collection, ReadOptions())
|
||||
|
||||
/**
|
||||
* Load data from ArangoDB into rdd
|
||||
*
|
||||
* @param sparkContext the sparkContext containing the ArangoDB configuration
|
||||
* @param collection the collection to load data from
|
||||
* @param additional read options
|
||||
*/
|
||||
def load[T: ClassTag](sparkContext: SparkContext, collection: String, options: ReadOptions): ArangoRdd[T] =
|
||||
new ArangoRdd[T](sparkContext, createReadOptions(options, sparkContext.getConf).copy(collection = collection))
|
||||
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
package cn.ac.iie.spark.partition
|
||||
|
||||
import org.apache.spark.Partition
|
||||
|
||||
class QueryArangoPartition(idx: Int, val offset: Long, val separate: Long) extends Partition{
|
||||
override def index: Int = idx
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
package cn.ac.iie.spark.rdd
|
||||
|
||||
import com.arangodb.Protocol
|
||||
import com.arangodb.entity.LoadBalancingStrategy
|
||||
|
||||
trait ArangoOptions {
|
||||
|
||||
def database: String = "_system"
|
||||
|
||||
def hosts: Option[String] = None
|
||||
|
||||
def user: Option[String] = None
|
||||
|
||||
def password: Option[String] = None
|
||||
|
||||
def useSsl: Option[Boolean] = None
|
||||
|
||||
def sslKeyStoreFile: Option[String] = None
|
||||
|
||||
def sslPassPhrase: Option[String] = None
|
||||
|
||||
def sslProtocol: Option[String] = None
|
||||
|
||||
def protocol: Option[Protocol] = None
|
||||
|
||||
def maxConnections: Option[Int] = None
|
||||
|
||||
def acquireHostList: Option[Boolean] = None
|
||||
|
||||
def acquireHostListInterval: Option[Int] = None
|
||||
|
||||
def loadBalancingStrategy: Option[LoadBalancingStrategy] = None
|
||||
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
package cn.ac.iie.spark.rdd
|
||||
|
||||
import scala.collection.JavaConverters.asScalaIteratorConverter
|
||||
import cn.ac.iie.config.ApplicationConfig
|
||||
import cn.ac.iie.service.update.UpdateDocument
|
||||
import cn.ac.iie.spark
|
||||
import cn.ac.iie.spark.partition.QueryArangoPartition
|
||||
import com.arangodb.ArangoCursor
|
||||
import org.apache.spark.{Partition, SparkContext, TaskContext}
|
||||
import org.apache.spark.rdd.RDD
|
||||
import org.slf4j.LoggerFactory
|
||||
|
||||
import scala.collection.mutable.ArrayBuffer
|
||||
import scala.reflect.ClassTag
|
||||
|
||||
class ArangoRdd[T: ClassTag](@transient override val sparkContext: SparkContext,
|
||||
val options: ReadOptions
|
||||
) extends RDD[T](sparkContext, Nil) {
|
||||
|
||||
private val LOG = LoggerFactory.getLogger(UpdateDocument.getClass)
|
||||
|
||||
override def compute(split: Partition, context: TaskContext): Iterator[T] = {
|
||||
|
||||
createCursor(split.asInstanceOf[QueryArangoPartition]).asScala
|
||||
}
|
||||
|
||||
override protected def getPartitions: Array[Partition] = {
|
||||
val partitions = ArrayBuffer[Partition]()
|
||||
val total = getCountTotal
|
||||
for (i <- 0 until ApplicationConfig.SPARK_SQL_SHUFFLE_PARTITIONS) {
|
||||
val partition = getPartition(i, total)
|
||||
partitions += partition
|
||||
}
|
||||
partitions.toArray
|
||||
}
|
||||
|
||||
private def createCursor(split: QueryArangoPartition)(implicit clazz: ClassTag[T]): ArangoCursor[T] = {
|
||||
|
||||
var arangoCursor:ArangoCursor[T] = null
|
||||
val arangoDB = spark.createArangoBuilder(options).build()
|
||||
try {
|
||||
val offset = split.offset
|
||||
val separate = split.separate
|
||||
val collection = options.collection
|
||||
val sql = s"FOR doc IN $collection limit $offset,$separate RETURN doc"
|
||||
LOG.info(sql)
|
||||
arangoCursor = arangoDB.db(options.database).query(sql,clazz.runtimeClass.asInstanceOf[Class[T]])
|
||||
}catch {
|
||||
case e: Exception => LOG.error("创建Cursor异常")
|
||||
}finally {
|
||||
arangoDB.shutdown()
|
||||
}
|
||||
arangoCursor
|
||||
}
|
||||
|
||||
override def repartition(numPartitions: Int)(implicit ord: Ordering[T]): RDD[T] = super.repartition(numPartitions)
|
||||
|
||||
private def getPartition(idx: Int, countTotal: Long): QueryArangoPartition = {
|
||||
val sepNum = countTotal / ApplicationConfig.THREAD_POOL_NUMBER + 1
|
||||
val offsetNum = idx * sepNum
|
||||
new QueryArangoPartition(idx, offsetNum, sepNum)
|
||||
}
|
||||
|
||||
override def count(): Long = getCountTotal
|
||||
|
||||
private def getCountTotal: Long = {
|
||||
val arangoDB = spark.createArangoBuilder(options).build()
|
||||
var cnt = 0L
|
||||
val sql = "RETURN LENGTH(" + options.collection + ")"
|
||||
LOG.info(sql)
|
||||
try {
|
||||
val longs = arangoDB.db(options.database).query(sql, classOf[Long])
|
||||
while (longs.hasNext) cnt = longs.next
|
||||
} catch {
|
||||
case e: Exception => LOG.error(sql + "执行异常")
|
||||
}finally {
|
||||
arangoDB.shutdown()
|
||||
}
|
||||
cnt
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
/*
|
||||
* DISCLAIMER
|
||||
*
|
||||
* Copyright 2016 ArangoDB GmbH, Cologne, Germany
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*
|
||||
* Copyright holder is ArangoDB GmbH, Cologne, Germany
|
||||
*
|
||||
* author Mark - mark at arangodb.com
|
||||
*/
|
||||
|
||||
package cn.ac.iie.spark.rdd
|
||||
|
||||
import cn.ac.iie.spark.partition.QueryArangoPartition
|
||||
import com.arangodb.Protocol
|
||||
import com.arangodb.entity.LoadBalancingStrategy
|
||||
|
||||
case class ReadOptions(override val database: String = "_system",
|
||||
val collection: String = null,
|
||||
partitioner: QueryArangoPartition = new QueryArangoPartition(0,0,0),
|
||||
override val hosts: Option[String] = None,
|
||||
override val user: Option[String] = None,
|
||||
override val password: Option[String] = None,
|
||||
override val useSsl: Option[Boolean] = None,
|
||||
override val sslKeyStoreFile: Option[String] = None,
|
||||
override val sslPassPhrase: Option[String] = None,
|
||||
override val sslProtocol: Option[String] = None,
|
||||
override val protocol: Option[Protocol] = None,
|
||||
override val maxConnections: Option[Int] = None,
|
||||
override val acquireHostList: Option[Boolean] = None,
|
||||
override val acquireHostListInterval: Option[Int] = None,
|
||||
override val loadBalancingStrategy: Option[LoadBalancingStrategy] = None) extends ArangoOptions {
|
||||
|
||||
def this() = this(database = "_system")
|
||||
|
||||
def database(database: String): ReadOptions = copy(database = database)
|
||||
|
||||
def collection(collection: String): ReadOptions = copy(collection = collection)
|
||||
|
||||
def hosts(hosts: String): ReadOptions = copy(hosts = Some(hosts))
|
||||
|
||||
def user(user: String): ReadOptions = copy(user = Some(user))
|
||||
|
||||
def password(password: String): ReadOptions = copy(password = Some(password))
|
||||
|
||||
def useSsl(useSsl: Boolean): ReadOptions = copy(useSsl = Some(useSsl))
|
||||
|
||||
def sslKeyStoreFile(sslKeyStoreFile: String): ReadOptions = copy(sslKeyStoreFile = Some(sslKeyStoreFile))
|
||||
|
||||
def sslPassPhrase(sslPassPhrase: String): ReadOptions = copy(sslPassPhrase = Some(sslPassPhrase))
|
||||
|
||||
def sslProtocol(sslProtocol: String): ReadOptions = copy(sslProtocol = Some(sslProtocol))
|
||||
|
||||
def protocol(protocol: Protocol): ReadOptions = copy(protocol = Some(protocol))
|
||||
|
||||
def maxConnections(maxConnections: Int): ReadOptions = copy(maxConnections = Some(maxConnections))
|
||||
|
||||
def acquireHostList(acquireHostList: Boolean): ReadOptions = copy(acquireHostList = Some(acquireHostList))
|
||||
|
||||
def acquireHostListInterval(acquireHostListInterval: Int): ReadOptions = copy(acquireHostListInterval = Some(acquireHostListInterval))
|
||||
|
||||
def loadBalancingStrategy(loadBalancingStrategy: LoadBalancingStrategy): ReadOptions = copy(loadBalancingStrategy = Some(loadBalancingStrategy))
|
||||
|
||||
def copy(database: String = database,
|
||||
collection: String = collection,
|
||||
partitioner: QueryArangoPartition = partitioner,
|
||||
hosts: Option[String] = hosts,
|
||||
user: Option[String] = user,
|
||||
password: Option[String] = password,
|
||||
useSsl: Option[Boolean] = useSsl,
|
||||
sslKeyStoreFile: Option[String] = sslKeyStoreFile,
|
||||
sslPassPhrase: Option[String] = sslPassPhrase,
|
||||
sslProtocol: Option[String] = sslProtocol,
|
||||
protocol: Option[Protocol] = protocol,
|
||||
maxConnections: Option[Int] = maxConnections,
|
||||
acquireHostList: Option[Boolean] = acquireHostList,
|
||||
acquireHostListInterval: Option[Int] = acquireHostListInterval,
|
||||
loadBalancingStrategy: Option[LoadBalancingStrategy] = loadBalancingStrategy): ReadOptions = {
|
||||
ReadOptions(database, collection, partitioner, hosts, user, password, useSsl, sslKeyStoreFile, sslPassPhrase, sslProtocol, protocol, maxConnections, acquireHostList, acquireHostListInterval, loadBalancingStrategy)
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,119 @@
|
||||
/*
|
||||
* DISCLAIMER
|
||||
*
|
||||
* Copyright 2016 ArangoDB GmbH, Cologne, Germany
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*
|
||||
* Copyright holder is ArangoDB GmbH, Cologne, Germany
|
||||
*
|
||||
* author Mark - mark at arangodb.com
|
||||
*/
|
||||
|
||||
package cn.ac.iie.spark.rdd
|
||||
|
||||
import javax.net.ssl.SSLContext
|
||||
import com.arangodb.Protocol
|
||||
import com.arangodb.entity.LoadBalancingStrategy
|
||||
|
||||
case class WriteOptions(override val database: String = "_system",
|
||||
val method: WriteOptions.Method = WriteOptions.INSERT,
|
||||
override val hosts: Option[String] = None,
|
||||
override val user: Option[String] = None,
|
||||
override val password: Option[String] = None,
|
||||
override val useSsl: Option[Boolean] = None,
|
||||
override val sslKeyStoreFile: Option[String] = None,
|
||||
override val sslPassPhrase: Option[String] = None,
|
||||
override val sslProtocol: Option[String] = None,
|
||||
override val protocol: Option[Protocol] = None,
|
||||
override val maxConnections: Option[Int] = None,
|
||||
override val acquireHostList: Option[Boolean] = None,
|
||||
override val acquireHostListInterval: Option[Int] = None,
|
||||
override val loadBalancingStrategy: Option[LoadBalancingStrategy] = None) extends ArangoOptions {
|
||||
import WriteOptions._
|
||||
|
||||
def this() = this(database = "_system")
|
||||
|
||||
def database(database: String): WriteOptions = copy(database = database)
|
||||
|
||||
def method(method: Method): WriteOptions = copy(method = method)
|
||||
|
||||
def hosts(hosts: String): WriteOptions = copy(hosts = Some(hosts))
|
||||
|
||||
def user(user: String): WriteOptions = copy(user = Some(user))
|
||||
|
||||
def password(password: String): WriteOptions = copy(password = Some(password))
|
||||
|
||||
def useSsl(useSsl: Boolean): WriteOptions = copy(useSsl = Some(useSsl))
|
||||
|
||||
def sslKeyStoreFile(sslKeyStoreFile: String): WriteOptions = copy(sslKeyStoreFile = Some(sslKeyStoreFile))
|
||||
|
||||
def sslPassPhrase(sslPassPhrase: String): WriteOptions = copy(sslPassPhrase = Some(sslPassPhrase))
|
||||
|
||||
def sslProtocol(sslProtocol: String): WriteOptions = copy(sslProtocol = Some(sslProtocol))
|
||||
|
||||
def protocol(protocol: Protocol): WriteOptions = copy(protocol = Some(protocol))
|
||||
|
||||
def maxConnections(maxConnections: Int): WriteOptions = copy(maxConnections = Some(maxConnections))
|
||||
|
||||
def acquireHostList(acquireHostList: Boolean): WriteOptions = copy(acquireHostList = Some(acquireHostList))
|
||||
|
||||
def acquireHostListInterval(acquireHostListInterval: Int): WriteOptions = copy(acquireHostListInterval = Some(acquireHostListInterval))
|
||||
|
||||
def loadBalancingStrategy(loadBalancingStrategy: LoadBalancingStrategy): WriteOptions = copy(loadBalancingStrategy = Some(loadBalancingStrategy))
|
||||
|
||||
def copy(database: String = database,
|
||||
method: Method = method,
|
||||
hosts: Option[String] = hosts,
|
||||
user: Option[String] = user,
|
||||
password: Option[String] = password,
|
||||
useSsl: Option[Boolean] = useSsl,
|
||||
sslKeyStoreFile: Option[String] = sslKeyStoreFile,
|
||||
sslPassPhrase: Option[String] = sslPassPhrase,
|
||||
sslProtocol: Option[String] = sslProtocol,
|
||||
protocol: Option[Protocol] = protocol,
|
||||
maxConnections: Option[Int] = maxConnections,
|
||||
acquireHostList: Option[Boolean] = acquireHostList,
|
||||
acquireHostListInterval: Option[Int] = acquireHostListInterval,
|
||||
loadBalancingStrategy: Option[LoadBalancingStrategy] = loadBalancingStrategy): WriteOptions = {
|
||||
WriteOptions(database, method, hosts, user, password, useSsl, sslKeyStoreFile, sslPassPhrase, sslProtocol, protocol, maxConnections, acquireHostList, acquireHostListInterval, loadBalancingStrategy)
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
object WriteOptions {
|
||||
|
||||
/**
|
||||
* method to save documents to arangodb
|
||||
*/
|
||||
sealed trait Method
|
||||
|
||||
/**
|
||||
* save documents by inserting
|
||||
* @see [[com.arangodb.ArangoCollection#insertDocuments(java.util.Collection)]]
|
||||
*/
|
||||
case object INSERT extends Method
|
||||
|
||||
/**
|
||||
* save documents by updating
|
||||
* @see [[com.arangodb.ArangoCollection#updateDocuments(java.util.Collection)]]
|
||||
*/
|
||||
case object UPDATE extends Method
|
||||
|
||||
/**
|
||||
* save documents by replacing
|
||||
* @see [[com.arangodb.ArangoCollection#replaceDocuments(java.util.Collection)]]
|
||||
*/
|
||||
case object REPLACE extends Method
|
||||
|
||||
}
|
||||
139
ip-learning-spark/src/main/scala/cn/ac/iie/spark/spark.scala
Normal file
139
ip-learning-spark/src/main/scala/cn/ac/iie/spark/spark.scala
Normal file
@@ -0,0 +1,139 @@
|
||||
/*
|
||||
* DISCLAIMER
|
||||
*
|
||||
* Copyright 2016 ArangoDB GmbH, Cologne, Germany
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*
|
||||
* Copyright holder is ArangoDB GmbH, Cologne, Germany
|
||||
*
|
||||
* author Mark - mark at arangodb.com
|
||||
*/
|
||||
|
||||
package cn.ac.iie
|
||||
|
||||
import java.io.FileInputStream
|
||||
import java.security.KeyStore
|
||||
|
||||
import cn.ac.iie.spark.rdd.{ArangoOptions, ReadOptions, WriteOptions}
|
||||
import com.arangodb.{ArangoDB, ArangoDBException, Protocol}
|
||||
import com.arangodb.entity.LoadBalancingStrategy
|
||||
import com.arangodb.velocypack.module.jdk8.VPackJdk8Module
|
||||
import com.arangodb.velocypack.module.scala.VPackScalaModule
|
||||
import javax.net.ssl.{KeyManagerFactory, SSLContext, TrustManagerFactory}
|
||||
import org.apache.spark.SparkConf
|
||||
|
||||
import scala.util.Try
|
||||
|
||||
package object spark {
|
||||
|
||||
val PropertyHosts = "arangodb.hosts"
|
||||
val PropertyUser = "arangodb.user"
|
||||
val PropertyPassword = "arangodb.password"
|
||||
val PropertyUseSsl = "arangodb.useSsl"
|
||||
val PropertySslKeyStoreFile = "arangodb.ssl.keyStoreFile"
|
||||
val PropertySslPassPhrase = "arangodb.ssl.passPhrase"
|
||||
val PropertySslProtocol = "arangodb.ssl.protocol"
|
||||
val PropertyProtocol = "arangodb.protocol"
|
||||
val PropertyMaxConnections = "arangodb.maxConnections"
|
||||
val PropertyAcquireHostList = "arangodb.acquireHostList"
|
||||
val PropertyAcquireHostListInterval = "arangodb.acquireHostListInterval"
|
||||
val PropertyLoadBalancingStrategy = "arangodb.loadBalancingStrategy"
|
||||
|
||||
private[spark] def createReadOptions(options: ReadOptions, sc: SparkConf): ReadOptions = {
|
||||
options.copy(
|
||||
hosts = options.hosts.orElse(some(sc.get(PropertyHosts, null))),
|
||||
user = options.user.orElse(some(sc.get(PropertyUser, null))),
|
||||
password = options.password.orElse(some(sc.get(PropertyPassword, null))),
|
||||
useSsl = options.useSsl.orElse(some(Try(sc.get(PropertyUseSsl, null).toBoolean).getOrElse(false))),
|
||||
sslKeyStoreFile = options.sslKeyStoreFile.orElse(some(sc.get(PropertySslKeyStoreFile, null))),
|
||||
sslPassPhrase = options.sslPassPhrase.orElse(some(sc.get(PropertySslPassPhrase, null))),
|
||||
sslProtocol = options.sslProtocol.orElse(some(sc.get(PropertySslProtocol, null))),
|
||||
protocol = options.protocol.orElse(some(Protocol.valueOf(sc.get(PropertyProtocol, "VST")))),
|
||||
maxConnections = options.maxConnections.orElse(some(Try(sc.get(PropertyMaxConnections, null).toInt).getOrElse(1))),
|
||||
acquireHostList = options.acquireHostList.orElse(some(Try(sc.get(PropertyAcquireHostList, null).toBoolean).getOrElse(false))),
|
||||
acquireHostListInterval = options.acquireHostListInterval.orElse(some(Try(sc.get(PropertyAcquireHostListInterval, null).toInt).getOrElse(60000))),
|
||||
loadBalancingStrategy = options.loadBalancingStrategy.orElse(some(LoadBalancingStrategy.valueOf(sc.get(PropertyLoadBalancingStrategy, "NONE")))))
|
||||
}
|
||||
|
||||
private[spark] def createWriteOptions(options: WriteOptions, sc: SparkConf): WriteOptions = {
|
||||
options.copy(
|
||||
hosts = options.hosts.orElse(some(sc.get(PropertyHosts, null))),
|
||||
user = options.user.orElse(some(sc.get(PropertyUser, null))),
|
||||
password = options.password.orElse(some(sc.get(PropertyPassword, null))),
|
||||
useSsl = options.useSsl.orElse(some(Try(sc.get(PropertyUseSsl, null).toBoolean).getOrElse(false))),
|
||||
sslKeyStoreFile = options.sslKeyStoreFile.orElse(some(sc.get(PropertySslKeyStoreFile, null))),
|
||||
sslPassPhrase = options.sslPassPhrase.orElse(some(sc.get(PropertySslPassPhrase, null))),
|
||||
sslProtocol = options.sslProtocol.orElse(some(sc.get(PropertySslProtocol, null))),
|
||||
protocol = options.protocol.orElse(some(Protocol.valueOf(sc.get(PropertyProtocol, "VST")))),
|
||||
maxConnections = options.maxConnections.orElse(some(Try(sc.get(PropertyMaxConnections, null).toInt).getOrElse(1))),
|
||||
acquireHostList = options.acquireHostList.orElse(some(Try(sc.get(PropertyAcquireHostList, null).toBoolean).getOrElse(false))),
|
||||
acquireHostListInterval = options.acquireHostListInterval.orElse(some(Try(sc.get(PropertyAcquireHostListInterval, null).toInt).getOrElse(60000))),
|
||||
loadBalancingStrategy = options.loadBalancingStrategy.orElse(some(LoadBalancingStrategy.valueOf(sc.get(PropertyLoadBalancingStrategy, "NONE")))))
|
||||
}
|
||||
|
||||
private[spark] def createArangoBuilder(options: ArangoOptions): ArangoDB.Builder = {
|
||||
val builder = new ArangoDB.Builder()
|
||||
builder.registerModules(new VPackJdk8Module, new VPackScalaModule)
|
||||
options.hosts.foreach { hosts(_).foreach(host => builder.host(host._1, host._2)) }
|
||||
options.user.foreach { builder.user(_) }
|
||||
options.password.foreach { builder.password(_) }
|
||||
options.useSsl.foreach { builder.useSsl(_) }
|
||||
if (options.sslKeyStoreFile.isDefined && options.sslPassPhrase.isDefined) {
|
||||
builder.sslContext(createSslContext(options.sslKeyStoreFile.get, options.sslPassPhrase.get, options.sslProtocol.getOrElse("TLS")))
|
||||
}
|
||||
options.protocol.foreach { builder.useProtocol(_) }
|
||||
options.maxConnections.foreach { builder.maxConnections(_) }
|
||||
options.acquireHostList.foreach { builder.acquireHostList(_) }
|
||||
options.acquireHostListInterval.foreach { builder.acquireHostListInterval(_) }
|
||||
options.loadBalancingStrategy.foreach { builder.loadBalancingStrategy(_) }
|
||||
builder
|
||||
}
|
||||
|
||||
private def createSslContext(keyStoreFile: String, passPhrase: String, protocol: String): SSLContext = {
|
||||
val ks = KeyStore.getInstance(KeyStore.getDefaultType());
|
||||
val kmf = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
|
||||
ks.load(new FileInputStream(keyStoreFile), passPhrase.toCharArray());
|
||||
kmf.init(ks, passPhrase.toCharArray());
|
||||
val tmf = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
|
||||
tmf.init(ks);
|
||||
val sc = SSLContext.getInstance(protocol);
|
||||
sc.init(kmf.getKeyManagers(), tmf.getTrustManagers(), null);
|
||||
sc
|
||||
}
|
||||
|
||||
private def some(value: String): Option[String] =
|
||||
if (value != null) Some(value) else None
|
||||
|
||||
private def some(value: Int): Option[Int] =
|
||||
Some(value)
|
||||
|
||||
private def some(value: Boolean): Option[Boolean] =
|
||||
Some(value)
|
||||
|
||||
private def some(value: Protocol): Option[Protocol] =
|
||||
Some(value)
|
||||
|
||||
private def some(value: LoadBalancingStrategy): Option[LoadBalancingStrategy] =
|
||||
Some(value)
|
||||
|
||||
private def hosts(hosts: String): List[(String, Int)] =
|
||||
hosts.split(",").map({ x =>
|
||||
val s = x.split(":")
|
||||
if (s.length != 2 || !s(1).matches("[0-9]+"))
|
||||
throw new ArangoDBException(s"Could not load property-value arangodb.hosts=${s}. Expected format ip:port,ip:port,...");
|
||||
else
|
||||
(s(0), s(1).toInt)
|
||||
}).toList
|
||||
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
/*
|
||||
* DISCLAIMER
|
||||
*
|
||||
* Copyright 2016 ArangoDB GmbH, Cologne, Germany
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*
|
||||
* Copyright holder is ArangoDB GmbH, Cologne, Germany
|
||||
*
|
||||
* author Mark - mark at arangodb.com
|
||||
*/
|
||||
|
||||
package cn.ac.iie.spark.vpack
|
||||
|
||||
import org.apache.spark.sql.Row
|
||||
import org.apache.spark.sql.types.{
|
||||
ArrayType,
|
||||
BooleanType,
|
||||
DataType,
|
||||
DateType,
|
||||
DecimalType,
|
||||
DoubleType,
|
||||
FloatType,
|
||||
IntegerType,
|
||||
LongType,
|
||||
MapType,
|
||||
NullType,
|
||||
ShortType,
|
||||
StringType,
|
||||
StructField,
|
||||
StructType,
|
||||
TimestampType
|
||||
}
|
||||
import com.arangodb.velocypack.VPackBuilder
|
||||
import com.arangodb.velocypack.VPackSlice
|
||||
import com.arangodb.velocypack.ValueType
|
||||
|
||||
private[spark] object VPackUtils {
|
||||
|
||||
def rowToVPack(row: Row): VPackSlice = {
|
||||
val builder = new VPackBuilder()
|
||||
if (row == null) {
|
||||
builder.add(ValueType.NULL)
|
||||
} else {
|
||||
builder.add(ValueType.OBJECT)
|
||||
row.schema.fields.zipWithIndex.foreach { addField(_, row, builder) }
|
||||
builder.close()
|
||||
}
|
||||
builder.slice()
|
||||
}
|
||||
|
||||
private def addField(field: (StructField, Int), row: Row, builder: VPackBuilder): Unit = {
|
||||
val name = field._1.name
|
||||
val index = field._2
|
||||
if (row.isNullAt(index)) {
|
||||
builder.add(name, ValueType.NULL)
|
||||
} else {
|
||||
field._1.dataType match {
|
||||
case BooleanType => builder.add(name, java.lang.Boolean.valueOf(row.getBoolean(index)))
|
||||
case DoubleType => builder.add(name, java.lang.Double.valueOf(row.getDouble(index)))
|
||||
case FloatType => builder.add(name, java.lang.Float.valueOf(row.getFloat(index)))
|
||||
case LongType => builder.add(name, java.lang.Long.valueOf(row.getLong(index)))
|
||||
case IntegerType => builder.add(name, java.lang.Integer.valueOf(row.getInt(index)))
|
||||
case ShortType => builder.add(name, java.lang.Short.valueOf(row.getShort(index)))
|
||||
case StringType => builder.add(name, java.lang.String.valueOf(row.getString(index)));
|
||||
case DateType => builder.add(name, row.getDate(index))
|
||||
case TimestampType => builder.add(name, row.getTimestamp(index))
|
||||
case t: DecimalType => builder.add(name, row.getDecimal(index))
|
||||
case t: MapType => {
|
||||
builder.add(name, ValueType.OBJECT)
|
||||
row.getMap[String, Any](index).foreach { case (name, value) => addValue(name, value, builder) }
|
||||
builder.close()
|
||||
}
|
||||
case t: ArrayType => {
|
||||
builder.add(name, ValueType.ARRAY)
|
||||
addValues(row, index, builder, t.elementType)
|
||||
builder.close()
|
||||
}
|
||||
case NullType => builder.add(name, ValueType.NULL)
|
||||
case struct: StructType => builder.add(name, rowToVPack(row.getStruct(index)))
|
||||
case _ => // TODO
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private def addValues(row: Row, index: Int, builder: VPackBuilder, itemType: DataType): Unit = {
|
||||
itemType match {
|
||||
case BooleanType =>
|
||||
row.getSeq[Boolean](index).foreach { value =>
|
||||
addValue(null, value, builder)
|
||||
}
|
||||
case DoubleType =>
|
||||
row.getSeq[Double](index).foreach { value =>
|
||||
addValue(null, value, builder)
|
||||
}
|
||||
case FloatType =>
|
||||
row.getSeq[Float](index).foreach { value =>
|
||||
addValue(null, value, builder)
|
||||
}
|
||||
case LongType =>
|
||||
row.getSeq[Long](index).foreach { value =>
|
||||
addValue(null, value, builder)
|
||||
}
|
||||
case IntegerType =>
|
||||
row.getSeq[Int](index).foreach { value =>
|
||||
addValue(null, value, builder)
|
||||
}
|
||||
case ShortType =>
|
||||
row.getSeq[Short](index).foreach { value =>
|
||||
addValue(null, value, builder)
|
||||
}
|
||||
case StringType =>
|
||||
row.getSeq[String](index).foreach { value =>
|
||||
addValue(null, value, builder)
|
||||
}
|
||||
case DateType =>
|
||||
row.getSeq[java.sql.Date](index).foreach { value =>
|
||||
addValue(null, value, builder)
|
||||
}
|
||||
case TimestampType =>
|
||||
row.getSeq[java.sql.Timestamp](index).foreach { value =>
|
||||
addValue(null, value, builder)
|
||||
}
|
||||
case s: StructType => {
|
||||
row.getSeq[Row](index).foreach { value =>
|
||||
builder.add(null, rowToVPack(value))
|
||||
}
|
||||
}
|
||||
case t: MapType => // TODO
|
||||
case t: ArrayType => // TODO
|
||||
case _ => // TODO
|
||||
}
|
||||
}
|
||||
|
||||
private def addValue(name: String, value: Any, builder: VPackBuilder): Unit = {
|
||||
value match {
|
||||
case value: Boolean => builder.add(name, java.lang.Boolean.valueOf(value))
|
||||
case value: Double => builder.add(name, java.lang.Double.valueOf(value))
|
||||
case value: Float => builder.add(name, java.lang.Float.valueOf(value))
|
||||
case value: Long => builder.add(name, java.lang.Long.valueOf(value))
|
||||
case value: Int => builder.add(name, java.lang.Integer.valueOf(value))
|
||||
case value: Short => builder.add(name, java.lang.Short.valueOf(value))
|
||||
case value: String => builder.add(name, java.lang.String.valueOf(value))
|
||||
case value: java.sql.Date => builder.add(name, value)
|
||||
case value: java.sql.Timestamp => builder.add(name, value)
|
||||
case _ => // TODO
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
@@ -17,6 +17,9 @@ object SparkSessionUtil {
|
||||
.config("spark.network.timeout", ApplicationConfig.SPARK_NETWORK_TIMEOUT)
|
||||
.config("spark.sql.shuffle.partitions", ApplicationConfig.SPARK_SQL_SHUFFLE_PARTITIONS)
|
||||
.config("spark.executor.memory", ApplicationConfig.SPARK_EXECUTOR_MEMORY)
|
||||
.config("arangodb.hosts", s"${ApplicationConfig.ARANGODB_HOST}:${ApplicationConfig.ARANGODB_PORT}")
|
||||
.config("arangodb.user", ApplicationConfig.ARANGODB_USER)
|
||||
.config("arangodb.password", ApplicationConfig.ARANGODB_PASSWORD)
|
||||
.master(ApplicationConfig.MASTER)
|
||||
.getOrCreate()
|
||||
LOG.warn("sparkession获取成功！！！")
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
package cn.ac.iie.spark
|
||||
|
||||
import cn.ac.iie.spark.rdd.ReadOptions
|
||||
import cn.ac.iie.utils.SparkSessionUtil
|
||||
import com.arangodb.entity.BaseDocument
|
||||
import org.apache.spark.rdd.RDD
|
||||
import org.apache.spark.storage.StorageLevel
|
||||
|
||||
object RDDTest {
|
||||
def main(args: Array[String]): Unit = {
|
||||
|
||||
val sparkContext = SparkSessionUtil.spark.sparkContext
|
||||
|
||||
println(sparkContext.getConf.get("arangodb.hosts"))
|
||||
|
||||
// val options = ReadOptions("iplearn_media_domain").copy(collection = "R_LOCATE_FQDN2IP")
|
||||
val options = ReadOptions("ip-learning-test-0")
|
||||
|
||||
val ipOptions = options.copy(collection = "IP")
|
||||
|
||||
val rdd = ArangoSpark.load[BaseDocument](sparkContext,"IP",options)
|
||||
|
||||
println(rdd.count())
|
||||
println(rdd.getNumPartitions)
|
||||
|
||||
val value: RDD[BaseDocument] = rdd.filter(doc => doc.getAttribute("CLIENT_SESSION_COUNT").asInstanceOf[Long] > 100).map(doc => {
|
||||
doc.addAttribute("abc", 1)
|
||||
doc
|
||||
})
|
||||
value.persist(StorageLevel.MEMORY_AND_DISK)
|
||||
|
||||
value.foreach(row => println(row.toString))
|
||||
println(value.count())
|
||||
|
||||
SparkSessionUtil.spark.close()
|
||||
System.exit(0)
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
Reference in New Issue
Block a user