java.nio.ByteBuffer Scala Examples

The following examples show how to use java.nio.ByteBuffer. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example.
Example 1
Source File: TsStreamingTest.scala    From spark-riak-connector   with Apache License 2.0 7 votes vote down vote up
package com.basho.riak.spark.streaming

import java.nio.ByteBuffer
import java.util.concurrent.{Callable, Executors, TimeUnit}

import com.basho.riak.spark._
import com.basho.riak.spark.rdd.RiakTSTests
import com.basho.riak.spark.rdd.timeseries.{AbstractTimeSeriesTest, TimeSeriesData}
import com.fasterxml.jackson.core.JsonParser
import com.fasterxml.jackson.databind.{DeserializationFeature, ObjectMapper, SerializationFeature}
import com.fasterxml.jackson.module.scala.DefaultScalaModule
import org.apache.spark.sql.Row
import org.junit.Assert._
import org.junit.experimental.categories.Category
import org.junit.{After, Before, Test}

@Category(Array(classOf[RiakTSTests]))
class TsStreamingTest extends AbstractTimeSeriesTest(false) with SparkStreamingFixture {

  protected final val executorService = Executors.newCachedThreadPool()
  private val dataSource = new SocketStreamingDataSource
  private var port = -1

  @Before
  def setUp(): Unit = {
    port = dataSource.start(client => {
      testData
        .map(tolerantMapper.writeValueAsString)
        .foreach(x => client.write(ByteBuffer.wrap(s"$x\n".getBytes)))
      logInfo(s"${testData.length} values were send to client")
    })
  }

  @After
  def tearDown(): Unit = {
    dataSource.stop()
  }

  @Test(timeout = 10 * 1000) // 10 seconds timeout
  def saveToRiak(): Unit = {
    executorService.submit(new Runnable {
      override def run(): Unit = {
        ssc.socketTextStream("localhost", port)
          .map(string => {
            val tsdata = new ObjectMapper()
              .configure(DeserializationFeature.FAIL_ON_NULL_FOR_PRIMITIVES, true)
              .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, true)
              .configure(JsonParser.Feature.ALLOW_SINGLE_QUOTES, true)
              .configure(JsonParser.Feature.ALLOW_UNQUOTED_FIELD_NAMES, true)
              .configure(SerializationFeature.WRITE_DATES_AS_TIMESTAMPS, false)
              .registerModule(DefaultScalaModule)
              .readValue(string, classOf[TimeSeriesData])
            Row(1, "f", tsdata.time, tsdata.user_id, tsdata.temperature_k)
          })
          .saveToRiakTS(bucketName)

        ssc.start()
        ssc.awaitTerminationOrTimeout(5 * 1000)
      }
    })

    val result = executorService.submit(new Callable[Array[Seq[Any]]] {
      override def call(): Array[Seq[Any]] = {
        var rdd = sc.riakTSTable[Row](bucketName)
          .sql(s"SELECT user_id, temperature_k FROM $bucketName $sqlWhereClause")
        var count = rdd.count()
        while (count < testData.length) {
          TimeUnit.SECONDS.sleep(2)

          rdd = sc.riakTSTable[Row](bucketName)
            .sql(s"SELECT user_id, temperature_k FROM $bucketName $sqlWhereClause")
          count = rdd.count()
        }
        rdd.collect().map(_.toSeq)
      }
    }).get()

    assertEquals(testData.length, result.length)
    assertEqualsUsingJSONIgnoreOrder(
      """
        |[
        |   ['bryce',305.37],
        |   ['bryce',300.12],
        |   ['bryce',295.95],
        |   ['ratman',362.121],
        |   ['ratman',3502.212]
        |]
      """.stripMargin, result)
  }
} 
Example 2
Source File: TrafficMonitorThread.scala    From shadowsocksr-android   with GNU General Public License v3.0 5 votes vote down vote up
package com.github.shadowsocks.utils

import java.io.{File, IOException}
import java.nio.{ByteBuffer, ByteOrder}
import java.util.concurrent.Executors

import android.content.Context
import android.net.{LocalServerSocket, LocalSocket, LocalSocketAddress}
import android.util.Log

class TrafficMonitorThread(context: Context) extends Thread {

  val TAG = "TrafficMonitorThread"
  lazy val PATH = context.getApplicationInfo.dataDir + "/stat_path"

  @volatile var serverSocket: LocalServerSocket = null
  @volatile var isRunning: Boolean = true

  def closeServerSocket() {
    if (serverSocket != null) {
      try {
        serverSocket.close()
      } catch {
        case _: Exception => // ignore
      }
      serverSocket = null
      }
  }

  def stopThread() {
    isRunning = false
    closeServerSocket()
  }

  override def run() {

    try {
      new File(PATH).delete()
    } catch {
      case _: Exception => // ignore
    }

    try {
      val localSocket = new LocalSocket
      localSocket.bind(new LocalSocketAddress(PATH, LocalSocketAddress.Namespace.FILESYSTEM))
      serverSocket = new LocalServerSocket(localSocket.getFileDescriptor)
    } catch {
      case e: IOException =>
        Log.e(TAG, "unable to bind", e)
        return
    }

    val pool = Executors.newFixedThreadPool(1)

    while (isRunning) {
      try {
        val socket = serverSocket.accept()

        pool.execute(() => {
          try {
            val input = socket.getInputStream
            val output = socket.getOutputStream

            val buffer = new Array[Byte](16)
            if (input.read(buffer) != 16) throw new IOException("Unexpected traffic stat length")
            val stat = ByteBuffer.wrap(buffer).order(ByteOrder.LITTLE_ENDIAN)
            TrafficMonitor.update(stat.getLong(0), stat.getLong(8))

            output.write(0)

            input.close()
            output.close()

          } catch {
            case e: Exception =>
              Log.e(TAG, "Error when recv traffic stat", e)
          }

          // close socket
          try {
            socket.close()
          } catch {
            case _: Exception => // ignore
          }

        })
      } catch {
        case e: IOException =>
          Log.e(TAG, "Error when accept socket", e)
          return
      }
    }
  }
} 
Example 3
Source File: JsMessageBuilder.scala    From iotchain   with MIT License 5 votes vote down vote up
package jbok.network.http

import java.nio.ByteBuffer

import cats.effect.{Async, IO}
import org.scalajs.dom.{Blob, Event, FileReader, UIEvent}

import scala.concurrent.Promise
import scala.scalajs.js.typedarray.TypedArrayBufferOps._
import scala.scalajs.js.typedarray._
import scala.scalajs.js.|

trait JsMessageBuilder[F[_], P] {
  import JsMessageBuilder._

  def responseType: String

  def pack(payload: P): Message

  def unpack(msg: Message): F[Option[P]]
}

object JsMessageBuilder {
  type Message = String | ArrayBuffer | Blob

  implicit def JsMessageBuilderString[F[_]](implicit F: Async[F]): JsMessageBuilder[F, String] = new JsMessageBuilder[F, String] {
    val responseType = ""

    def pack(payload: String): Message = payload

    def unpack(msg: Message): F[Option[String]] = (msg: Any) match {
      case s: String => F.pure(Some(s))
      case b: Blob   => readBlob[F, String, String](_.readAsText(b))(identity)
      case _         => F.pure(None)
    }
  }

  implicit def JsMessageBuilderByteBuffer[F[_]](implicit F: Async[F]): JsMessageBuilder[F, ByteBuffer] = new JsMessageBuilder[F, ByteBuffer] {
    val responseType = "arraybuffer"

    def pack(payload: ByteBuffer): Message = payload.arrayBuffer.slice(payload.position, payload.limit)

    def unpack(msg: Message): F[Option[ByteBuffer]] = (msg: Any) match {
      case a: ArrayBuffer => F.pure(Option(TypedArrayBuffer.wrap(a)))
      case b: Blob        => readBlob[F, ArrayBuffer, ByteBuffer](_.readAsArrayBuffer(b))(TypedArrayBuffer.wrap(_))
      case _              => F.pure(None)
    }
  }

  private def readBlob[F[_], R, W](doRead: FileReader => Unit)(conv: R => W)(implicit F: Async[F]): F[Option[W]] = {
    val promise = Promise[Option[W]]()
    val reader  = new FileReader
    reader.onload = (_: UIEvent) => {
      val s = reader.result.asInstanceOf[R]
      promise.success(Option(conv(s)))
    }
    reader.onerror = (_: Event) => {
      promise.success(None)
    }
    doRead(reader)
    F.liftIO(IO.fromFuture(IO(promise.future)))
  }
} 
Example 4
Source File: scodec.scala    From iotchain   with MIT License 5 votes vote down vote up
package jbok.codec.impl

import java.nio.ByteBuffer

import _root_.scodec.bits.BitVector
import jbok.codec._

object scodec {
  final case class DeserializeException(msg: String) extends Exception(msg)

  implicit def scodecCodec[T: _root_.scodec.Codec]: Codec2[T, ByteBuffer] =
    new Encoder[T, ByteBuffer] with Decoder[T, ByteBuffer] {
      override def encode(arg: T): ByteBuffer = _root_.scodec.Codec[T].encode(arg).require.toByteBuffer
      override def decode(arg: ByteBuffer): Either[Throwable, T] =
        _root_.scodec.Codec[T].decode(BitVector(arg)).toEither match {
          case Right(result) => Right(result.value)
          case Left(err)     => Left(DeserializeException(err.messageWithContext))
        }
    }
} 
Example 5
Source File: ByteUtils.scala    From iotchain   with MIT License 5 votes vote down vote up
package jbok.common

import java.nio.{ByteBuffer, ByteOrder}

import scodec.bits.ByteVector

object ByteUtils {
  def or(arrays: ByteVector*): ByteVector = {
    require(arrays.map(_.length).distinct.length <= 1, "All the arrays should have the same length")
    require(arrays.nonEmpty, "There should be one or more arrays")

    val zeroes = ByteVector.fill(arrays.headOption.map(_.length).getOrElse(0))(0.toByte)
    arrays.foldLeft[ByteVector](zeroes) {
      case (acc, cur) => acc or cur
    }
  }

  def and(arrays: ByteVector*): ByteVector = {
    require(arrays.map(_.length).distinct.length <= 1, "All the arrays should have the same length")
    require(arrays.nonEmpty, "There should be one or more arrays")

    val ones = ByteVector.fill(arrays.headOption.map(_.length).getOrElse(0))(0xFF.toByte)
    arrays.foldLeft[ByteVector](ones) {
      case (acc, cur) => acc and cur
    }
  }

  def bytesToInts(bytes: Array[Byte]): Array[Int] =
    bytes.grouped(4).map(getIntFromWord).toArray

  def intsToBytes(input: Array[Int]): Array[Byte] =
    input.flatMap { i =>
      Array(
        (i & 0xFF).toByte,
        ((i >> 8) & 0xFF).toByte,
        ((i >> 16) & 0xFF).toByte,
        ((i >> 24) & 0xFF).toByte
      )
    }

  def getIntFromWord(arr: Array[Byte]): Int =
    ByteBuffer.wrap(arr, 0, 4).order(ByteOrder.LITTLE_ENDIAN).getInt
} 
Example 6
Source File: MesosTaskLaunchData.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.scheduler.cluster.mesos

import java.nio.ByteBuffer

import org.apache.mesos.protobuf.ByteString

import org.apache.spark.internal.Logging


private[spark] case class MesosTaskLaunchData(
  serializedTask: ByteBuffer,
  attemptNumber: Int) extends Logging {

  def toByteString: ByteString = {
    val dataBuffer = ByteBuffer.allocate(4 + serializedTask.limit)
    dataBuffer.putInt(attemptNumber)
    dataBuffer.put(serializedTask)
    dataBuffer.rewind
    logDebug(s"ByteBuffer size: [${dataBuffer.remaining}]")
    ByteString.copyFrom(dataBuffer)
  }
}

private[spark] object MesosTaskLaunchData extends Logging {
  def fromByteString(byteString: ByteString): MesosTaskLaunchData = {
    val byteBuffer = byteString.asReadOnlyByteBuffer()
    logDebug(s"ByteBuffer size: [${byteBuffer.remaining}]")
    val attemptNumber = byteBuffer.getInt // updates the position by 4 bytes
    val serializedTask = byteBuffer.slice() // subsequence starting at the current position
    MesosTaskLaunchData(serializedTask, attemptNumber)
  }
} 
Example 7
Source File: MesosTaskLaunchDataSuite.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.scheduler.cluster.mesos

import java.nio.ByteBuffer

import org.apache.spark.SparkFunSuite

class MesosTaskLaunchDataSuite extends SparkFunSuite {
  test("serialize and deserialize data must be same") {
    val serializedTask = ByteBuffer.allocate(40)
    (Range(100, 110).map(serializedTask.putInt(_)))
    serializedTask.rewind
    val attemptNumber = 100
    val byteString = MesosTaskLaunchData(serializedTask, attemptNumber).toByteString
    serializedTask.rewind
    val mesosTaskLaunchData = MesosTaskLaunchData.fromByteString(byteString)
    assert(mesosTaskLaunchData.attemptNumber == attemptNumber)
    assert(mesosTaskLaunchData.serializedTask.equals(serializedTask))
  }
} 
Example 8
Source File: KPLBasedKinesisTestUtils.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.streaming.kinesis

import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets

import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer

import com.amazonaws.services.kinesis.producer.{KinesisProducer => KPLProducer, KinesisProducerConfiguration, UserRecordResult}
import com.google.common.util.concurrent.{FutureCallback, Futures}

private[kinesis] class KPLBasedKinesisTestUtils extends KinesisTestUtils {
  override protected def getProducer(aggregate: Boolean): KinesisDataGenerator = {
    if (!aggregate) {
      new SimpleDataGenerator(kinesisClient)
    } else {
      new KPLDataGenerator(regionName)
    }
  }
}


private[kinesis] class KPLDataGenerator(regionName: String) extends KinesisDataGenerator {

  private lazy val producer: KPLProducer = {
    val conf = new KinesisProducerConfiguration()
      .setRecordMaxBufferedTime(1000)
      .setMaxConnections(1)
      .setRegion(regionName)
      .setMetricsLevel("none")

    new KPLProducer(conf)
  }

  override def sendData(streamName: String, data: Seq[Int]): Map[String, Seq[(Int, String)]] = {
    val shardIdToSeqNumbers = new mutable.HashMap[String, ArrayBuffer[(Int, String)]]()
    data.foreach { num =>
      val str = num.toString
      val data = ByteBuffer.wrap(str.getBytes(StandardCharsets.UTF_8))
      val future = producer.addUserRecord(streamName, str, data)
      val kinesisCallBack = new FutureCallback[UserRecordResult]() {
        override def onFailure(t: Throwable): Unit = {} // do nothing

        override def onSuccess(result: UserRecordResult): Unit = {
          val shardId = result.getShardId
          val seqNumber = result.getSequenceNumber()
          val sentSeqNumbers = shardIdToSeqNumbers.getOrElseUpdate(shardId,
            new ArrayBuffer[(Int, String)]())
          sentSeqNumbers += ((num, seqNumber))
        }
      }
      Futures.addCallback(future, kinesisCallBack)
    }
    producer.flushSync()
    shardIdToSeqNumbers.toMap
  }
} 
Example 9
Source File: FlumeInputDStream.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.streaming.flume

import java.io.{Externalizable, ObjectInput, ObjectOutput}
import java.net.InetSocketAddress
import java.nio.ByteBuffer
import java.util.concurrent.Executors

import scala.collection.JavaConverters._
import scala.reflect.ClassTag

import org.apache.avro.ipc.NettyServer
import org.apache.avro.ipc.specific.SpecificResponder
import org.apache.flume.source.avro.{AvroFlumeEvent, AvroSourceProtocol, Status}
import org.jboss.netty.channel.{ChannelPipeline, ChannelPipelineFactory, Channels}
import org.jboss.netty.channel.socket.nio.NioServerSocketChannelFactory
import org.jboss.netty.handler.codec.compression._

import org.apache.spark.internal.Logging
import org.apache.spark.storage.StorageLevel
import org.apache.spark.streaming.StreamingContext
import org.apache.spark.streaming.dstream._
import org.apache.spark.streaming.receiver.Receiver
import org.apache.spark.util.Utils

private[streaming]
class FlumeInputDStream[T: ClassTag](
  _ssc: StreamingContext,
  host: String,
  port: Int,
  storageLevel: StorageLevel,
  enableDecompression: Boolean
) extends ReceiverInputDStream[SparkFlumeEvent](_ssc) {

  override def getReceiver(): Receiver[SparkFlumeEvent] = {
    new FlumeReceiver(host, port, storageLevel, enableDecompression)
  }
}


  private[streaming]
  class CompressionChannelPipelineFactory extends ChannelPipelineFactory {
    def getPipeline(): ChannelPipeline = {
      val pipeline = Channels.pipeline()
      val encoder = new ZlibEncoder(6)
      pipeline.addFirst("deflater", encoder)
      pipeline.addFirst("inflater", new ZlibDecoder())
      pipeline
    }
  }
} 
Example 10
Source File: FlumeTestUtils.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.streaming.flume

import java.net.{InetSocketAddress, ServerSocket}
import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets
import java.util.{List => JList}
import java.util.Collections

import scala.collection.JavaConverters._

import org.apache.avro.ipc.NettyTransceiver
import org.apache.avro.ipc.specific.SpecificRequestor
import org.apache.commons.lang3.RandomUtils
import org.apache.flume.source.avro
import org.apache.flume.source.avro.{AvroFlumeEvent, AvroSourceProtocol}
import org.jboss.netty.channel.ChannelPipeline
import org.jboss.netty.channel.socket.SocketChannel
import org.jboss.netty.channel.socket.nio.NioClientSocketChannelFactory
import org.jboss.netty.handler.codec.compression.{ZlibDecoder, ZlibEncoder}

import org.apache.spark.util.Utils
import org.apache.spark.SparkConf


  private class CompressionChannelFactory(compressionLevel: Int)
    extends NioClientSocketChannelFactory {

    override def newChannel(pipeline: ChannelPipeline): SocketChannel = {
      val encoder = new ZlibEncoder(compressionLevel)
      pipeline.addFirst("deflater", encoder)
      pipeline.addFirst("inflater", new ZlibDecoder())
      super.newChannel(pipeline)
    }
  }

} 
Example 11
Source File: NullableColumnAccessor.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.sql.execution.columnar

import java.nio.{ByteBuffer, ByteOrder}

import org.apache.spark.sql.catalyst.InternalRow

private[columnar] trait NullableColumnAccessor extends ColumnAccessor {
  private var nullsBuffer: ByteBuffer = _
  private var nullCount: Int = _
  private var seenNulls: Int = 0

  private var nextNullIndex: Int = _
  private var pos: Int = 0

  abstract override protected def initialize(): Unit = {
    nullsBuffer = underlyingBuffer.duplicate().order(ByteOrder.nativeOrder())
    nullCount = ByteBufferHelper.getInt(nullsBuffer)
    nextNullIndex = if (nullCount > 0) ByteBufferHelper.getInt(nullsBuffer) else -1
    pos = 0

    underlyingBuffer.position(underlyingBuffer.position + 4 + nullCount * 4)
    super.initialize()
  }

  abstract override def extractTo(row: InternalRow, ordinal: Int): Unit = {
    if (pos == nextNullIndex) {
      seenNulls += 1

      if (seenNulls < nullCount) {
        nextNullIndex = ByteBufferHelper.getInt(nullsBuffer)
      }

      row.setNullAt(ordinal)
    } else {
      super.extractTo(row, ordinal)
    }

    pos += 1
  }

  abstract override def hasNext: Boolean = seenNulls < nullCount || super.hasNext
} 
Example 12
Source File: NullableColumnBuilder.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.sql.execution.columnar

import java.nio.{ByteBuffer, ByteOrder}

import org.apache.spark.sql.catalyst.InternalRow


private[columnar] trait NullableColumnBuilder extends ColumnBuilder {
  protected var nulls: ByteBuffer = _
  protected var nullCount: Int = _
  private var pos: Int = _

  abstract override def initialize(
      initialSize: Int,
      columnName: String,
      useCompression: Boolean): Unit = {

    nulls = ByteBuffer.allocate(1024)
    nulls.order(ByteOrder.nativeOrder())
    pos = 0
    nullCount = 0
    super.initialize(initialSize, columnName, useCompression)
  }

  abstract override def appendFrom(row: InternalRow, ordinal: Int): Unit = {
    columnStats.gatherStats(row, ordinal)
    if (row.isNullAt(ordinal)) {
      nulls = ColumnBuilder.ensureFreeSpace(nulls, 4)
      nulls.putInt(pos)
      nullCount += 1
    } else {
      super.appendFrom(row, ordinal)
    }
    pos += 1
  }

  abstract override def build(): ByteBuffer = {
    val nonNulls = super.build()
    val nullDataLen = nulls.position()

    nulls.limit(nullDataLen)
    nulls.rewind()

    val buffer = ByteBuffer
      .allocate(4 + nullDataLen + nonNulls.remaining())
      .order(ByteOrder.nativeOrder())
      .putInt(nullCount)
      .put(nulls)
      .put(nonNulls)

    buffer.rewind()
    buffer
  }

  protected def buildNonNulls(): ByteBuffer = {
    nulls.limit(nulls.position()).rewind()
    super.build()
  }
} 
Example 13
Source File: CompressibleColumnBuilder.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.sql.execution.columnar.compression

import java.nio.{ByteBuffer, ByteOrder}

import org.apache.spark.internal.Logging
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.execution.columnar.{ColumnBuilder, NativeColumnBuilder}
import org.apache.spark.sql.types.AtomicType
import org.apache.spark.unsafe.Platform


private[columnar] trait CompressibleColumnBuilder[T <: AtomicType]
  extends ColumnBuilder with Logging {

  this: NativeColumnBuilder[T] with WithCompressionSchemes =>

  var compressionEncoders: Seq[Encoder[T]] = _

  abstract override def initialize(
      initialSize: Int,
      columnName: String,
      useCompression: Boolean): Unit = {

    compressionEncoders =
      if (useCompression) {
        schemes.filter(_.supports(columnType)).map(_.encoder[T](columnType))
      } else {
        Seq(PassThrough.encoder(columnType))
      }
    super.initialize(initialSize, columnName, useCompression)
  }

  // The various compression schemes, while saving memory use, cause all of the data within
  // the row to become unaligned, thus causing crashes.  Until a way of fixing the compression
  // is found to also allow aligned accesses this must be disabled for SPARC.

  protected def isWorthCompressing(encoder: Encoder[T]) = {
    CompressibleColumnBuilder.unaligned && encoder.compressionRatio < 0.8
  }

  private def gatherCompressibilityStats(row: InternalRow, ordinal: Int): Unit = {
    compressionEncoders.foreach(_.gatherCompressibilityStats(row, ordinal))
  }

  abstract override def appendFrom(row: InternalRow, ordinal: Int): Unit = {
    super.appendFrom(row, ordinal)
    if (!row.isNullAt(ordinal)) {
      gatherCompressibilityStats(row, ordinal)
    }
  }

  override def build(): ByteBuffer = {
    val nonNullBuffer = buildNonNulls()
    val encoder: Encoder[T] = {
      val candidate = compressionEncoders.minBy(_.compressionRatio)
      if (isWorthCompressing(candidate)) candidate else PassThrough.encoder(columnType)
    }

    // Header = null count + null positions
    val headerSize = 4 + nulls.limit()
    val compressedSize = if (encoder.compressedSize == 0) {
      nonNullBuffer.remaining()
    } else {
      encoder.compressedSize
    }

    val compressedBuffer = ByteBuffer
      // Reserves 4 bytes for compression scheme ID
      .allocate(headerSize + 4 + compressedSize)
      .order(ByteOrder.nativeOrder)
      // Write the header
      .putInt(nullCount)
      .put(nulls)

    logDebug(s"Compressor for [$columnName]: $encoder, ratio: ${encoder.compressionRatio}")
    encoder.compress(nonNullBuffer, compressedBuffer)
  }
}

private[columnar] object CompressibleColumnBuilder {
  val unaligned = Platform.unaligned()
} 
Example 14
Source File: NullableColumnAccessorSuite.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.sql.execution.columnar

import java.nio.ByteBuffer

import org.apache.spark.SparkFunSuite
import org.apache.spark.sql.catalyst.CatalystTypeConverters
import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeProjection}
import org.apache.spark.sql.types._

class TestNullableColumnAccessor[JvmType](
    buffer: ByteBuffer,
    columnType: ColumnType[JvmType])
  extends BasicColumnAccessor(buffer, columnType)
  with NullableColumnAccessor

object TestNullableColumnAccessor {
  def apply[JvmType](buffer: ByteBuffer, columnType: ColumnType[JvmType])
    : TestNullableColumnAccessor[JvmType] = {
    new TestNullableColumnAccessor(buffer, columnType)
  }
}

class NullableColumnAccessorSuite extends SparkFunSuite {
  import org.apache.spark.sql.execution.columnar.ColumnarTestUtils._

  Seq(
    NULL, BOOLEAN, BYTE, SHORT, INT, LONG, FLOAT, DOUBLE,
    STRING, BINARY, COMPACT_DECIMAL(15, 10), LARGE_DECIMAL(20, 10),
    STRUCT(StructType(StructField("a", StringType) :: Nil)),
    ARRAY(ArrayType(IntegerType)), MAP(MapType(IntegerType, StringType)))
    .foreach {
    testNullableColumnAccessor(_)
  }

  def testNullableColumnAccessor[JvmType](
      columnType: ColumnType[JvmType]): Unit = {

    val typeName = columnType.getClass.getSimpleName.stripSuffix("$")
    val nullRow = makeNullRow(1)

    test(s"Nullable $typeName column accessor: empty column") {
      val builder = TestNullableColumnBuilder(columnType)
      val accessor = TestNullableColumnAccessor(builder.build(), columnType)
      assert(!accessor.hasNext)
    }

    test(s"Nullable $typeName column accessor: access null values") {
      val builder = TestNullableColumnBuilder(columnType)
      val randomRow = makeRandomRow(columnType)
      val proj = UnsafeProjection.create(Array[DataType](columnType.dataType))

      (0 until 4).foreach { _ =>
        builder.appendFrom(proj(randomRow), 0)
        builder.appendFrom(proj(nullRow), 0)
      }

      val accessor = TestNullableColumnAccessor(builder.build(), columnType)
      val row = new GenericInternalRow(1)
      val converter = CatalystTypeConverters.createToScalaConverter(columnType.dataType)

      (0 until 4).foreach { _ =>
        assert(accessor.hasNext)
        accessor.extractTo(row, 0)
        assert(converter(row.get(0, columnType.dataType))
          === converter(randomRow.get(0, columnType.dataType)))

        assert(accessor.hasNext)
        accessor.extractTo(row, 0)
        assert(row.isNullAt(0))
      }

      assert(!accessor.hasNext)
    }
  }
} 
Example 15
Source File: RawTextSender.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.streaming.util

import java.io.{ByteArrayOutputStream, IOException}
import java.net.ServerSocket
import java.nio.ByteBuffer

import scala.io.Source

import org.apache.spark.SparkConf
import org.apache.spark.internal.Logging
import org.apache.spark.serializer.KryoSerializer
import org.apache.spark.util.IntParam


private[streaming]
object RawTextSender extends Logging {
  def main(args: Array[String]) {
    if (args.length != 4) {
      // scalastyle:off println
      System.err.println("Usage: RawTextSender <port> <file> <blockSize> <bytesPerSec>")
      // scalastyle:on println
      System.exit(1)
    }
    // Parse the arguments using a pattern match
    val Array(IntParam(port), file, IntParam(blockSize), IntParam(bytesPerSec)) = args

    // Repeat the input data multiple times to fill in a buffer
    val lines = Source.fromFile(file).getLines().toArray
    val bufferStream = new ByteArrayOutputStream(blockSize + 1000)
    val ser = new KryoSerializer(new SparkConf()).newInstance()
    val serStream = ser.serializeStream(bufferStream)
    var i = 0
    while (bufferStream.size < blockSize) {
      serStream.writeObject(lines(i))
      i = (i + 1) % lines.length
    }
    val array = bufferStream.toByteArray

    val countBuf = ByteBuffer.wrap(new Array[Byte](4))
    countBuf.putInt(array.length)
    countBuf.flip()

    val serverSocket = new ServerSocket(port)
    logInfo("Listening on port " + port)

    while (true) {
      val socket = serverSocket.accept()
      logInfo("Got a new connection")
      val out = new RateLimitedOutputStream(socket.getOutputStream, bytesPerSec)
      try {
        while (true) {
          out.write(countBuf.array)
          out.write(array)
        }
      } catch {
        case e: IOException =>
          logError("Client disconnected")
      } finally {
        socket.close()
      }
    }
  }
} 
Example 16
Source File: FileBasedWriteAheadLogRandomReader.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.streaming.util

import java.io.Closeable
import java.nio.ByteBuffer

import org.apache.hadoop.conf.Configuration


private[streaming] class FileBasedWriteAheadLogRandomReader(path: String, conf: Configuration)
  extends Closeable {

  private val instream = HdfsUtils.getInputStream(path, conf)
  private var closed = (instream == null) // the file may be deleted as we're opening the stream

  def read(segment: FileBasedWriteAheadLogSegment): ByteBuffer = synchronized {
    assertOpen()
    instream.seek(segment.offset)
    val nextLength = instream.readInt()
    HdfsUtils.checkState(nextLength == segment.length,
      s"Expected message length to be ${segment.length}, but was $nextLength")
    val buffer = new Array[Byte](nextLength)
    instream.readFully(buffer)
    ByteBuffer.wrap(buffer)
  }

  override def close(): Unit = synchronized {
    closed = true
    instream.close()
  }

  private def assertOpen() {
    HdfsUtils.checkState(!closed, "Stream is closed. Create a new Reader to read from the file.")
  }
} 
Example 17
Source File: FileBasedWriteAheadLogWriter.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.streaming.util

import java.io._
import java.nio.ByteBuffer

import org.apache.hadoop.conf.Configuration

import org.apache.spark.util.Utils


  def write(data: ByteBuffer): FileBasedWriteAheadLogSegment = synchronized {
    assertOpen()
    data.rewind() // Rewind to ensure all data in the buffer is retrieved
    val lengthToWrite = data.remaining()
    val segment = new FileBasedWriteAheadLogSegment(path, nextOffset, lengthToWrite)
    stream.writeInt(lengthToWrite)
    Utils.writeByteBuffer(data, stream: OutputStream)
    flush()
    nextOffset = stream.getPos()
    segment
  }

  override def close(): Unit = synchronized {
    closed = true
    stream.close()
  }

  private def flush() {
    stream.hflush()
    // Useful for local file system where hflush/sync does not work (HADOOP-7844)
    stream.getWrappedStream.flush()
  }

  private def assertOpen() {
    HdfsUtils.checkState(!closed, "Stream is closed. Create a new Writer to write to file.")
  }
} 
Example 18
Source File: FileBasedWriteAheadLogReader.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.streaming.util

import java.io.{Closeable, EOFException, IOException}
import java.nio.ByteBuffer

import org.apache.hadoop.conf.Configuration

import org.apache.spark.internal.Logging


private[streaming] class FileBasedWriteAheadLogReader(path: String, conf: Configuration)
  extends Iterator[ByteBuffer] with Closeable with Logging {

  private val instream = HdfsUtils.getInputStream(path, conf)
  private var closed = (instream == null) // the file may be deleted as we're opening the stream
  private var nextItem: Option[ByteBuffer] = None

  override def hasNext: Boolean = synchronized {
    if (closed) {
      return false
    }

    if (nextItem.isDefined) { // handle the case where hasNext is called without calling next
      true
    } else {
      try {
        val length = instream.readInt()
        val buffer = new Array[Byte](length)
        instream.readFully(buffer)
        nextItem = Some(ByteBuffer.wrap(buffer))
        logTrace("Read next item " + nextItem.get)
        true
      } catch {
        case e: EOFException =>
          logDebug("Error reading next item, EOF reached", e)
          close()
          false
        case e: IOException =>
          logWarning("Error while trying to read data. If the file was deleted, " +
            "this should be okay.", e)
          close()
          if (HdfsUtils.checkFileExists(path, conf)) {
            // If file exists, this could be a legitimate error
            throw e
          } else {
            // File was deleted. This can occur when the daemon cleanup thread takes time to
            // delete the file during recovery.
            false
          }

        case e: Exception =>
          logWarning("Error while trying to read data from HDFS.", e)
          close()
          throw e
      }
    }
  }

  override def next(): ByteBuffer = synchronized {
    val data = nextItem.getOrElse {
      close()
      throw new IllegalStateException(
        "next called without calling hasNext or after hasNext returned false")
    }
    nextItem = None // Ensure the next hasNext call loads new data.
    data
  }

  override def close(): Unit = synchronized {
    if (!closed) {
      instream.close()
    }
    closed = true
  }
} 
Example 19
Source File: BlockTransferService.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.network

import java.io.Closeable
import java.nio.ByteBuffer

import scala.concurrent.{Future, Promise}
import scala.concurrent.duration.Duration
import scala.reflect.ClassTag

import org.apache.spark.internal.Logging
import org.apache.spark.network.buffer.{ManagedBuffer, NioManagedBuffer}
import org.apache.spark.network.shuffle.{BlockFetchingListener, ShuffleClient}
import org.apache.spark.scheduler.MapStatus
import org.apache.spark.storage.{BlockId, StorageLevel}
import org.apache.spark.util.ThreadUtils

private[spark]
abstract class BlockTransferService extends ShuffleClient with Closeable with Logging {

  
  def uploadBlockSync(
      hostname: String,
      port: Int,
      execId: String,
      blockId: BlockId,
      blockData: ManagedBuffer,
      level: StorageLevel,
      classTag: ClassTag[_]): Unit = {
    val future = uploadBlock(hostname, port, execId, blockId, blockData, level, classTag)
    ThreadUtils.awaitResult(future, Duration.Inf)
  }
} 
Example 20
Source File: NettyBlockRpcServer.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.network.netty

import java.nio.ByteBuffer

import scala.collection.JavaConverters._
import scala.language.existentials
import scala.reflect.ClassTag

import org.apache.spark.internal.Logging
import org.apache.spark.network.BlockDataManager
import org.apache.spark.network.buffer.{ManagedBuffer, NioManagedBuffer}
import org.apache.spark.network.client.{RpcResponseCallback, TransportClient}
import org.apache.spark.network.server.{OneForOneStreamManager, RpcHandler, StreamManager}
import org.apache.spark.network.shuffle.protocol.{BlockTransferMessage, MapOutputReady, OpenBlocks, StreamHandle, UploadBlock}
import org.apache.spark.scheduler.MapStatus
import org.apache.spark.serializer.Serializer
import org.apache.spark.storage.{BlockId, StorageLevel}


class NettyBlockRpcServer(
    appId: String,
    serializer: Serializer,
    blockManager: BlockDataManager)
  extends RpcHandler with Logging {

  private val streamManager = new OneForOneStreamManager()

  override def receive(
      client: TransportClient,
      rpcMessage: ByteBuffer,
      responseContext: RpcResponseCallback): Unit = {
    val message = BlockTransferMessage.Decoder.fromByteBuffer(rpcMessage)
    logTrace(s"Received request: $message")

    message match {
      case openBlocks: OpenBlocks =>
        val blocks: Seq[ManagedBuffer] =
          openBlocks.blockIds.map(BlockId.apply).map(blockManager.getBlockData)
        val streamId = streamManager.registerStream(appId, blocks.iterator.asJava)
        logTrace(s"Registered streamId $streamId with ${blocks.size} buffers")
        responseContext.onSuccess(new StreamHandle(streamId, blocks.size).toByteBuffer)

      case uploadBlock: UploadBlock =>
        // StorageLevel and ClassTag are serialized as bytes using our JavaSerializer.
        val (level: StorageLevel, classTag: ClassTag[_]) = {
          serializer
            .newInstance()
            .deserialize(ByteBuffer.wrap(uploadBlock.metadata))
            .asInstanceOf[(StorageLevel, ClassTag[_])]
        }
        val data = new NioManagedBuffer(ByteBuffer.wrap(uploadBlock.blockData))
        val blockId = BlockId(uploadBlock.blockId)
        blockManager.putBlockData(blockId, data, level, classTag)
        responseContext.onSuccess(ByteBuffer.allocate(0))

      case mapOutputReady: MapOutputReady =>
        val mapStatus: MapStatus =
          serializer.newInstance().deserialize(ByteBuffer.wrap(mapOutputReady.serializedMapStatus))
        blockManager.mapOutputReady(
          mapOutputReady.shuffleId, mapOutputReady.mapId, mapOutputReady.numReduces, mapStatus)
    }
  }

  override def getStreamManager(): StreamManager = streamManager
} 
Example 21
Source File: ZooKeeperPersistenceEngine.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.deploy.master

import java.nio.ByteBuffer

import scala.collection.JavaConverters._
import scala.reflect.ClassTag

import org.apache.curator.framework.CuratorFramework
import org.apache.zookeeper.CreateMode

import org.apache.spark.SparkConf
import org.apache.spark.deploy.SparkCuratorUtil
import org.apache.spark.internal.Logging
import org.apache.spark.serializer.Serializer


private[master] class ZooKeeperPersistenceEngine(conf: SparkConf, val serializer: Serializer)
  extends PersistenceEngine
  with Logging {

  private val WORKING_DIR = conf.get("spark.deploy.zookeeper.dir", "/spark") + "/master_status"
  private val zk: CuratorFramework = SparkCuratorUtil.newClient(conf)

  SparkCuratorUtil.mkdir(zk, WORKING_DIR)


  override def persist(name: String, obj: Object): Unit = {
    serializeIntoFile(WORKING_DIR + "/" + name, obj)
  }

  override def unpersist(name: String): Unit = {
    zk.delete().forPath(WORKING_DIR + "/" + name)
  }

  override def read[T: ClassTag](prefix: String): Seq[T] = {
    zk.getChildren.forPath(WORKING_DIR).asScala
      .filter(_.startsWith(prefix)).flatMap(deserializeFromFile[T])
  }

  override def close() {
    zk.close()
  }

  private def serializeIntoFile(path: String, value: AnyRef) {
    val serialized = serializer.newInstance().serialize(value)
    val bytes = new Array[Byte](serialized.remaining())
    serialized.get(bytes)
    zk.create().withMode(CreateMode.PERSISTENT).forPath(path, bytes)
  }

  private def deserializeFromFile[T](filename: String)(implicit m: ClassTag[T]): Option[T] = {
    val fileData = zk.getData().forPath(WORKING_DIR + "/" + filename)
    try {
      Some(serializer.newInstance().deserialize[T](ByteBuffer.wrap(fileData)))
    } catch {
      case e: Exception =>
        logWarning("Exception while reading persisted file, deleting", e)
        zk.delete().forPath(WORKING_DIR + "/" + filename)
        None
    }
  }
} 
Example 22
Source File: JavaSerializer.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.serializer

import java.io._
import java.nio.ByteBuffer

import scala.reflect.ClassTag

import org.apache.spark.SparkConf
import org.apache.spark.annotation.DeveloperApi
import org.apache.spark.util.{ByteBufferInputStream, ByteBufferOutputStream, Utils}

private[spark] class JavaSerializationStream(
    out: OutputStream, counterReset: Int, extraDebugInfo: Boolean)
  extends SerializationStream {
  private val objOut = new ObjectOutputStream(out)
  private var counter = 0

  
@DeveloperApi
class JavaSerializer(conf: SparkConf) extends Serializer with Externalizable {
  private var counterReset = conf.getInt("spark.serializer.objectStreamReset", 100)
  private var extraDebugInfo = conf.getBoolean("spark.serializer.extraDebugInfo", true)

  protected def this() = this(new SparkConf())  // For deserialization only

  override def newInstance(): SerializerInstance = {
    val classLoader = defaultClassLoader.getOrElse(Thread.currentThread.getContextClassLoader)
    new JavaSerializerInstance(counterReset, extraDebugInfo, classLoader)
  }

  override def writeExternal(out: ObjectOutput): Unit = Utils.tryOrIOException {
    out.writeInt(counterReset)
    out.writeBoolean(extraDebugInfo)
  }

  override def readExternal(in: ObjectInput): Unit = Utils.tryOrIOException {
    counterReset = in.readInt()
    extraDebugInfo = in.readBoolean()
  }
} 
Example 23
Source File: LocalSchedulerBackend.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.scheduler.local

import java.io.File
import java.net.URL
import java.nio.ByteBuffer

import org.apache.spark.{SparkConf, SparkContext, SparkEnv, TaskState}
import org.apache.spark.TaskState.TaskState
import org.apache.spark.executor.{Executor, ExecutorBackend}
import org.apache.spark.internal.Logging
import org.apache.spark.launcher.{LauncherBackend, SparkAppHandle}
import org.apache.spark.rpc.{RpcCallContext, RpcEndpointRef, RpcEnv, ThreadSafeRpcEndpoint}
import org.apache.spark.scheduler._
import org.apache.spark.scheduler.cluster.ExecutorInfo

private case class ReviveOffers()

private case class StatusUpdate(taskId: Long, state: TaskState, serializedData: ByteBuffer)

private case class KillTask(taskId: Long, interruptThread: Boolean)

private case class StopExecutor()


  def getUserClasspath(conf: SparkConf): Seq[URL] = {
    val userClassPathStr = conf.getOption("spark.executor.extraClassPath")
    userClassPathStr.map(_.split(File.pathSeparator)).toSeq.flatten.map(new File(_).toURI.toURL)
  }

  launcherBackend.connect()

  override def start() {
    val rpcEnv = SparkEnv.get.rpcEnv
    val executorEndpoint = new LocalEndpoint(rpcEnv, userClassPath, scheduler, this, totalCores)
    localEndpoint = rpcEnv.setupEndpoint("LocalSchedulerBackendEndpoint", executorEndpoint)
    listenerBus.post(SparkListenerExecutorAdded(
      System.currentTimeMillis,
      executorEndpoint.localExecutorId,
      new ExecutorInfo(executorEndpoint.localExecutorHostname, totalCores, Map.empty)))
    launcherBackend.setAppId(appId)
    launcherBackend.setState(SparkAppHandle.State.RUNNING)
  }

  override def stop() {
    stop(SparkAppHandle.State.FINISHED)
  }

  override def reviveOffers() {
    localEndpoint.send(ReviveOffers)
  }

  override def defaultParallelism(): Int =
    scheduler.conf.getInt("spark.default.parallelism", totalCores)

  override def killTask(taskId: Long, executorId: String, interruptThread: Boolean) {
    localEndpoint.send(KillTask(taskId, interruptThread))
  }

  override def statusUpdate(taskId: Long, state: TaskState, serializedData: ByteBuffer) {
    localEndpoint.send(StatusUpdate(taskId, state, serializedData))
  }

  override def applicationId(): String = appId

  private def stop(finalState: SparkAppHandle.State): Unit = {
    localEndpoint.ask(StopExecutor)
    try {
      launcherBackend.setState(finalState)
    } finally {
      launcherBackend.close()
    }
  }

} 
Example 24
Source File: ShuffleMapTask.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.scheduler

import java.lang.management.ManagementFactory
import java.nio.ByteBuffer
import java.util.Properties

import scala.language.existentials

import org.apache.spark._
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.executor.TaskMetrics
import org.apache.spark.internal.Logging
import org.apache.spark.rdd.RDD
import org.apache.spark.shuffle.ShuffleWriter
import org.apache.spark.storage.BlockManagerId


  def this(partitionId: Int) {
    this(0, 0, null, new Partition { override def index: Int = 0 }, null, new Properties, null)
  }

  @transient private val preferredLocs: Seq[TaskLocation] = {
    if (locs == null) Nil else locs.toSet.toSeq
  }

  var rdd: RDD[_] = null
  var dep: ShuffleDependency[_, _, _] = null

  override def prepTask(): Unit = {
    // Deserialize the RDD using the broadcast variable.
    val threadMXBean = ManagementFactory.getThreadMXBean
    val deserializeStartTime = System.currentTimeMillis()
    val deserializeStartCpuTime = if (threadMXBean.isCurrentThreadCpuTimeSupported) {
      threadMXBean.getCurrentThreadCpuTime
    } else 0L
    val ser = SparkEnv.get.closureSerializer.newInstance()
    val (_rdd, _dep) = ser.deserialize[(RDD[_], ShuffleDependency[_, _, _])](
      ByteBuffer.wrap(taskBinary.value), Thread.currentThread.getContextClassLoader)
     rdd = _rdd
     dep = _dep
    _executorDeserializeTime = System.currentTimeMillis() - deserializeStartTime
    _executorDeserializeCpuTime = if (threadMXBean.isCurrentThreadCpuTimeSupported) {
      threadMXBean.getCurrentThreadCpuTime - deserializeStartCpuTime
    } else 0L
  }

  override def runTask(context: TaskContext): MapStatus = {
    if (dep == null || rdd == null) {
      prepTask()
    }

    var writer: ShuffleWriter[Any, Any] = null
    try {
      val manager = SparkEnv.get.shuffleManager
      writer = manager.getWriter[Any, Any](dep.shuffleHandle, partitionId, context)
      writer.write(rdd.iterator(partition, context).asInstanceOf[Iterator[_ <: Product2[Any, Any]]])
      val status = writer.stop(success = true).get
      FutureTaskNotifier.taskCompleted(status, partitionId, dep.shuffleId,
        dep.partitioner.numPartitions, nextStageLocs, metrics.shuffleWriteMetrics, false)
      status
    } catch {
      case e: Exception =>
        try {
          if (writer != null) {
            writer.stop(success = false)
          }
        } catch {
          case e: Exception =>
            log.debug("Could not stop writer", e)
        }
        throw e
    }
  }

  override def preferredLocations: Seq[TaskLocation] = preferredLocs

  override def toString: String = "ShuffleMapTask(%d, %d)".format(stageId, partitionId)
}

object ShuffleMapTask {

  def apply(
      stageId: Int,
      stageAttemptId: Int,
      partition: Partition,
      properties: Properties,
      internalAccumulatorsSer: Array[Byte],
      isFutureTask: Boolean,
      rdd: RDD[_],
      dep: ShuffleDependency[_, _, _],
      nextStageLocs: Option[Seq[BlockManagerId]]): ShuffleMapTask = {

    val smt = new ShuffleMapTask(stageId, stageAttemptId, null, partition, null,
      properties, internalAccumulatorsSer, isFutureTask, nextStageLocs)

    smt.rdd = rdd
    smt.dep = dep
    smt
  }
} 
Example 25
Source File: TaskDescription.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.scheduler

import java.nio.ByteBuffer

import scala.collection.mutable
import scala.collection.mutable.HashSet
import scala.util.control.NonFatal

import org.apache.spark._
import org.apache.spark.internal.Logging
import org.apache.spark.serializer.SerializerInstance
import org.apache.spark.util.SerializableBuffer


private[spark] class TaskDescription(
    val taskId: Long,
    val attemptNumber: Int,
    val executorId: String,
    val name: String,
    val index: Int,    // Index within this task's TaskSet
    val isFutureTask: Boolean,
    @transient private val _task: Task[_],
    @transient private val _addedFiles: mutable.Map[String, Long],
    @transient private val _addedJars: mutable.Map[String, Long],
    @transient private val _ser: SerializerInstance)
  extends Serializable with Logging {

  // Because ByteBuffers are not serializable, wrap the task in a SerializableBuffer
  private var buffer: SerializableBuffer = _

  def prepareSerializedTask(): Unit = {
    if (_task != null) {
      val serializedTask: ByteBuffer = try {
        Task.serializeWithDependencies(_task, _addedFiles, _addedJars, _ser)
      } catch {
        // If the task cannot be serialized, then there is not point in re-attempting
        // the task as it will always fail. So just abort the task set.
        case NonFatal(e) =>
          val msg = s"Failed to serialize the task $taskId, not attempting to retry it."
          logError(msg, e)
          // FIXME(shivaram): We dont have a handle to the taskSet here to abort it.
          throw new TaskNotSerializableException(e)
      }
      if (serializedTask.limit > TaskSetManager.TASK_SIZE_TO_WARN_KB * 1024) {
        logWarning(s"Stage ${_task.stageId} contains a task of very large size " +
          s"(${serializedTask.limit / 1024} KB). The maximum recommended task size is " +
          s"${TaskSetManager.TASK_SIZE_TO_WARN_KB} KB.")
      }
      buffer = new SerializableBuffer(serializedTask)
    } else {
      buffer = new SerializableBuffer(ByteBuffer.allocate(0))
    }
  }

  def serializedTask: ByteBuffer = buffer.value

  override def toString: String = "TaskDescription(TID=%d, index=%d)".format(taskId, index)
} 
Example 26
Source File: TaskResult.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.scheduler

import java.io._
import java.nio.ByteBuffer

import scala.collection.mutable.ArrayBuffer

import org.apache.spark.SparkEnv
import org.apache.spark.serializer.SerializerInstance
import org.apache.spark.storage.BlockId
import org.apache.spark.util.{AccumulatorV2, Utils}

// Task result. Also contains updates to accumulator variables.
private[spark] sealed trait TaskResult[T]


  def value(resultSer: SerializerInstance = null): T = {
    if (valueObjectDeserialized) {
      valueObject
    } else {
      // This should not run when holding a lock because it may cost dozens of seconds for a large
      // value
      val ser = if (resultSer == null) SparkEnv.get.serializer.newInstance() else resultSer
      valueObject = ser.deserialize(valueBytes)
      valueObjectDeserialized = true
      valueObject
    }
  }
} 
Example 27
Source File: ResultTask.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.scheduler

import java.io._
import java.lang.management.ManagementFactory
import java.nio.ByteBuffer
import java.util.Properties

import org.apache.spark._
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.executor.TaskMetrics
import org.apache.spark.rdd.RDD


private[spark] class ResultTask[T, U](
    stageId: Int,
    stageAttemptId: Int,
    taskBinary: Broadcast[Array[Byte]],
    partition: Partition,
    locs: Seq[TaskLocation],
    val outputId: Int,
    localProperties: Properties,
    serializedTaskMetrics: Array[Byte] =
      SparkEnv.get.closureSerializer.newInstance().serialize(TaskMetrics.registered).array(),
    isFutureTask: Boolean = false,
    depShuffleIds: Option[Seq[Seq[Int]]] = None,
    depShuffleNumMaps: Option[Seq[Int]] = None,
    jobId: Option[Int] = None,
    appId: Option[String] = None,
    appAttemptId: Option[String] = None)
  extends Task[U](stageId, stageAttemptId, partition.index,
    serializedTaskMetrics, localProperties, isFutureTask, depShuffleIds, depShuffleNumMaps,
    jobId, appId, appAttemptId)
  with Serializable {

  var rdd: RDD[T] = null
  var func: (TaskContext, Iterator[T]) => U = null

  @transient private[this] val preferredLocs: Seq[TaskLocation] = {
    if (locs == null) Nil else locs.toSet.toSeq
  }

  override def prepTask(): Unit = {
    // Deserialize the RDD and the func using the broadcast variables.
    val threadMXBean = ManagementFactory.getThreadMXBean
    val deserializeStartTime = System.currentTimeMillis()
    val deserializeStartCpuTime = if (threadMXBean.isCurrentThreadCpuTimeSupported) {
      threadMXBean.getCurrentThreadCpuTime
    } else 0L
    val ser = SparkEnv.get.closureSerializer.newInstance()
    val (_rdd, _func) = ser.deserialize[(RDD[T], (TaskContext, Iterator[T]) => U)](
      ByteBuffer.wrap(taskBinary.value), Thread.currentThread.getContextClassLoader)
    rdd = _rdd
    func = _func
    _executorDeserializeTime = System.currentTimeMillis() - deserializeStartTime
    _executorDeserializeCpuTime = if (threadMXBean.isCurrentThreadCpuTimeSupported) {
      threadMXBean.getCurrentThreadCpuTime - deserializeStartCpuTime
    } else 0L
  }

  override def runTask(context: TaskContext): U = {
    // Deserialize the RDD and the func using the broadcast variables.
    if (func == null || rdd == null) {
      prepTask()
    }
    func(context, rdd.iterator(partition, context))
  }

  // This is only callable on the driver side.
  override def preferredLocations: Seq[TaskLocation] = preferredLocs

  override def toString: String = "ResultTask(" + stageId + ", " + partitionId + ")"
}

object ResultTask {

  def apply[T, U](
      stageId: Int,
      stageAttemptId: Int,
      partition: Partition,
      outputId: Int,
      localProperties: Properties,
      internalAccumulatorsSer: Array[Byte],
      isFutureTask: Boolean,
      rdd: RDD[T],
      func: (TaskContext, Iterator[T]) => U): ResultTask[T, U] = {
    val rt = new ResultTask[T, U](stageId, stageAttemptId, null, partition, Seq.empty, outputId,
      localProperties, internalAccumulatorsSer, isFutureTask)
    rt.rdd = rdd
    rt.func = func
    rt
  }

} 
Example 28
Source File: BatchResultTask.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.scheduler

import java.io._
import java.nio.ByteBuffer
import java.util.Properties

import scala.reflect.ClassTag

import org.apache.spark._
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.rdd.RDD

private[spark] class BatchResultTask[T, U: ClassTag](
    stageId: Int,
    stageAttemptId: Int,
    taskBinaries: Broadcast[Array[Byte]],
    val partitions: Array[Partition],
    partitionId: Int,
    @transient private val locs: Seq[TaskLocation],
    val outputId: Int,
    localProperties: Properties,
    internalAccumulatorsSer: Array[Byte],
    isFutureTask: Boolean,
    depShuffleIds: Option[Seq[Seq[Int]]] = None,
    depShuffleNumMaps: Option[Seq[Int]] = None,
    jobId: Option[Int] = None,
    appId: Option[String] = None,
    appAttemptId: Option[String] = None)
  extends Task[Array[U]](stageId, stageAttemptId, partitionId,
      internalAccumulatorsSer, localProperties, isFutureTask, depShuffleIds, depShuffleNumMaps,
      jobId, appId, appAttemptId)
  with BatchTask
  with Serializable {

  @transient private[this] val preferredLocs: Seq[TaskLocation] = {
    if (locs == null) Nil else locs.toSet.toSeq
  }

  var rdds: Array[RDD[T]] = null

  var funcs: Array[(TaskContext, Iterator[T]) => U] = null

  override def prepTask(): Unit = {
    // Deserialize the RDD and the func using the broadcast variables.
    val ser = SparkEnv.get.closureSerializer.newInstance()
    val (rddI, funcI) =
      ser.deserialize[(Array[RDD[T]], Array[(TaskContext, Iterator[T]) => U])](
        ByteBuffer.wrap(taskBinaries.value), Thread.currentThread.getContextClassLoader)
    rdds = rddI
    funcs = funcI
  }

  // Called on the executor side to get a smaller tasks out
  def getTasks(): Seq[Task[Any]] = {
    if (rdds == null) {
      prepTask()
    }

    (0 until partitions.length).map { i =>
      val r = ResultTask(stageId, stageAttemptId, partitions(i), outputId, localProperties,
        internalAccumulatorsSer, isFutureTask, rdds(i), funcs(i))
      r.epoch = epoch
      r
    }.map(_.asInstanceOf[Task[Any]])
  }

  override def runTask(context: TaskContext): Array[U] = {
    throw new RuntimeException("BatchResultTasks should not be run!")
  }

  // This is only callable on the driver side.
  override def preferredLocations: Seq[TaskLocation] = preferredLocs

  override def toString: String = "BatchResultTask(" + stageId + ", " + partitionId + ")"
} 
Example 29
Source File: BatchShuffleMapTask.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.scheduler

import java.io._
import java.nio.ByteBuffer
import java.util.Properties

import scala.reflect.ClassTag

import org.apache.spark._
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.internal.Logging
import org.apache.spark.rdd.RDD
import org.apache.spark.storage.BlockManagerId

private[spark] class BatchShuffleMapTask(
    stageId: Int,
    stageAttemptId: Int,
    taskBinaries: Broadcast[Array[Byte]],
    partitions: Array[Partition],
    partitionId: Int,
    @transient private var locs: Seq[TaskLocation],
    internalAccumulatorsSer: Array[Byte],
    localProperties: Properties,
    isFutureTask: Boolean,
    nextStageLocs: Option[Seq[BlockManagerId]] = None,
    depShuffleIds: Option[Seq[Seq[Int]]] = None,
    depShuffleNumMaps: Option[Seq[Int]] = None,
    jobId: Option[Int] = None,
    appId: Option[String] = None,
    appAttemptId: Option[String] = None)
  extends Task[Array[MapStatus]](stageId, stageAttemptId, partitionId,
    internalAccumulatorsSer, localProperties, isFutureTask, depShuffleIds, depShuffleNumMaps,
    jobId, appId, appAttemptId)
  with BatchTask
  with Logging {

  @transient private val preferredLocs: Seq[TaskLocation] = {
    if (locs == null) Nil else locs.toSet.toSeq
  }

  var rdds: Array[RDD[_]] = null
  var deps: Array[ShuffleDependency[_, _, _]] = null

  override def prepTask(): Unit = {
    // Deserialize the RDD using the broadcast variable.
    val ser = SparkEnv.get.closureSerializer.newInstance()
    val (rddI, depI) = ser.deserialize[(Array[RDD[_]], Array[ShuffleDependency[_, _, _]])](
      ByteBuffer.wrap(taskBinaries.value), Thread.currentThread.getContextClassLoader)
    rdds = rddI
    deps = depI
  }

  def getTasks(): Seq[Task[Any]] = {
    if (deps == null || rdds == null) {
      prepTask()
    }

    (0 until partitions.length).map { i =>
      val s = ShuffleMapTask(stageId, stageAttemptId, partitions(i), localProperties,
        internalAccumulatorsSer, isFutureTask, rdds(i), deps(i), nextStageLocs)
      s.epoch = epoch
      s
    }.map(_.asInstanceOf[Task[Any]])
  }

  override def runTask(context: TaskContext): Array[MapStatus] = {
    throw new RuntimeException("BatchShuffleMapTasks should not be run!")
  }

  override def preferredLocations: Seq[TaskLocation] = preferredLocs

  override def toString: String = "BatchShuffleMapTask(%d, %d)".format(stageId, partitionId)
} 
Example 30
Source File: ByteBufferOutputStream.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.util

import java.io.ByteArrayOutputStream
import java.nio.ByteBuffer


private[spark] class ByteBufferOutputStream(capacity: Int) extends ByteArrayOutputStream(capacity) {

  def this() = this(32)

  def getCount(): Int = count

  private[this] var closed: Boolean = false

  override def write(b: Int): Unit = {
    require(!closed, "cannot write to a closed ByteBufferOutputStream")
    super.write(b)
  }

  override def write(b: Array[Byte], off: Int, len: Int): Unit = {
    require(!closed, "cannot write to a closed ByteBufferOutputStream")
    super.write(b, off, len)
  }

  override def reset(): Unit = {
    require(!closed, "cannot reset a closed ByteBufferOutputStream")
    super.reset()
  }

  override def close(): Unit = {
    if (!closed) {
      super.close()
      closed = true
    }
  }

  def toByteBuffer: ByteBuffer = {
    require(closed, "can only call toByteBuffer() after ByteBufferOutputStream has been closed")
    ByteBuffer.wrap(buf, 0, count)
  }
} 
Example 31
Source File: SerializableBuffer.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.util

import java.io.{EOFException, IOException, ObjectInputStream, ObjectOutputStream}
import java.nio.ByteBuffer
import java.nio.channels.Channels


private[spark]
class SerializableBuffer(@transient var buffer: ByteBuffer) extends Serializable {
  def value: ByteBuffer = buffer

  private def readObject(in: ObjectInputStream): Unit = Utils.tryOrIOException {
    val length = in.readInt()
    buffer = ByteBuffer.allocate(length)
    var amountRead = 0
    val channel = Channels.newChannel(in)
    while (amountRead < length) {
      val ret = channel.read(buffer)
      if (ret == -1) {
        throw new EOFException("End of file before fully reading buffer")
      }
      amountRead += ret
    }
    buffer.rewind() // Allow us to read it later
  }

  private def writeObject(out: ObjectOutputStream): Unit = Utils.tryOrIOException {
    out.writeInt(buffer.limit())
    if (Channels.newChannel(out).write(buffer) != buffer.limit()) {
      throw new IOException("Could not fully write buffer to output stream")
    }
    buffer.rewind() // Allow us to write it again later
  }
} 
Example 32
Source File: ChunkedByteBufferOutputStream.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.util.io

import java.io.OutputStream
import java.nio.ByteBuffer

import scala.collection.mutable.ArrayBuffer

import org.apache.spark.storage.StorageUtils


  private[this] var position = chunkSize
  private[this] var _size = 0
  private[this] var closed: Boolean = false

  def size: Long = _size

  override def close(): Unit = {
    if (!closed) {
      super.close()
      closed = true
    }
  }

  override def write(b: Int): Unit = {
    require(!closed, "cannot write to a closed ChunkedByteBufferOutputStream")
    allocateNewChunkIfNeeded()
    chunks(lastChunkIndex).put(b.toByte)
    position += 1
    _size += 1
  }

  override def write(bytes: Array[Byte], off: Int, len: Int): Unit = {
    require(!closed, "cannot write to a closed ChunkedByteBufferOutputStream")
    var written = 0
    while (written < len) {
      allocateNewChunkIfNeeded()
      val thisBatch = math.min(chunkSize - position, len - written)
      chunks(lastChunkIndex).put(bytes, written + off, thisBatch)
      written += thisBatch
      position += thisBatch
    }
    _size += len
  }

  @inline
  private def allocateNewChunkIfNeeded(): Unit = {
    if (position == chunkSize) {
      chunks += allocator(chunkSize)
      lastChunkIndex += 1
      position = 0
    }
  }

  def toChunkedByteBuffer: ChunkedByteBuffer = {
    require(closed, "cannot call toChunkedByteBuffer() unless close() has been called")
    require(!toChunkedByteBufferWasCalled, "toChunkedByteBuffer() can only be called once")
    toChunkedByteBufferWasCalled = true
    if (lastChunkIndex == -1) {
      new ChunkedByteBuffer(Array.empty[ByteBuffer])
    } else {
      // Copy the first n-1 chunks to the output, and then create an array that fits the last chunk.
      // An alternative would have been returning an array of ByteBuffers, with the last buffer
      // bounded to only the last chunk's position. However, given our use case in Spark (to put
      // the chunks in block manager), only limiting the view bound of the buffer would still
      // require the block manager to store the whole chunk.
      val ret = new Array[ByteBuffer](chunks.size)
      for (i <- 0 until chunks.size - 1) {
        ret(i) = chunks(i)
        ret(i).flip()
      }
      if (position == chunkSize) {
        ret(lastChunkIndex) = chunks(lastChunkIndex)
        ret(lastChunkIndex).flip()
      } else {
        ret(lastChunkIndex) = allocator(position)
        chunks(lastChunkIndex).flip()
        ret(lastChunkIndex).put(chunks(lastChunkIndex))
        ret(lastChunkIndex).flip()
        StorageUtils.dispose(chunks(lastChunkIndex))
      }
      new ChunkedByteBuffer(ret)
    }
  }
} 
Example 33
Source File: DiskStore.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.storage

import java.io.{FileOutputStream, IOException, RandomAccessFile}
import java.nio.ByteBuffer
import java.nio.channels.FileChannel.MapMode

import com.google.common.io.Closeables

import org.apache.spark.SparkConf
import org.apache.spark.internal.Logging
import org.apache.spark.util.Utils
import org.apache.spark.util.io.ChunkedByteBuffer


  def put(blockId: BlockId)(writeFunc: FileOutputStream => Unit): Unit = {
    if (contains(blockId)) {
      throw new IllegalStateException(s"Block $blockId is already present in the disk store")
    }
    logDebug(s"Attempting to put block $blockId")
    val startTime = System.currentTimeMillis
    val file = diskManager.getFile(blockId)
    val fileOutputStream = new FileOutputStream(file)
    var threwException: Boolean = true
    try {
      writeFunc(fileOutputStream)
      threwException = false
    } finally {
      try {
        Closeables.close(fileOutputStream, threwException)
      } finally {
         if (threwException) {
          remove(blockId)
        }
      }
    }
    val finishTime = System.currentTimeMillis
    logDebug("Block %s stored as %s file on disk in %d ms".format(
      file.getName,
      Utils.bytesToString(file.length()),
      finishTime - startTime))
  }

  def putBytes(blockId: BlockId, bytes: ChunkedByteBuffer): Unit = {
    put(blockId) { fileOutputStream =>
      val channel = fileOutputStream.getChannel
      Utils.tryWithSafeFinally {
        bytes.writeFully(channel)
      } {
        channel.close()
      }
    }
  }

  def getBytes(blockId: BlockId): ChunkedByteBuffer = {
    val file = diskManager.getFile(blockId.name)
    val channel = new RandomAccessFile(file, "r").getChannel
    Utils.tryWithSafeFinally {
      // For small files, directly read rather than memory map
      if (file.length < minMemoryMapBytes) {
        val buf = ByteBuffer.allocate(file.length.toInt)
        channel.position(0)
        while (buf.remaining() != 0) {
          if (channel.read(buf) == -1) {
            throw new IOException("Reached EOF before filling buffer\n" +
              s"offset=0\nfile=${file.getAbsolutePath}\nbuf.remaining=${buf.remaining}")
          }
        }
        buf.flip()
        new ChunkedByteBuffer(buf)
      } else {
        new ChunkedByteBuffer(channel.map(MapMode.READ_ONLY, 0, file.length))
      }
    } {
      channel.close()
    }
  }

  def remove(blockId: BlockId): Boolean = {
    val file = diskManager.getFile(blockId.name)
    if (file.exists()) {
      val ret = file.delete()
      if (!ret) {
        logWarning(s"Error deleting ${file.getPath()}")
      }
      ret
    } else {
      false
    }
  }

  def contains(blockId: BlockId): Boolean = {
    val file = diskManager.getFile(blockId.name)
    file.exists()
  }
} 
Example 34
Source File: NettyRpcHandlerSuite.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.rpc.netty

import java.net.InetSocketAddress
import java.nio.ByteBuffer

import io.netty.channel.Channel
import org.mockito.Matchers._
import org.mockito.Mockito._

import org.apache.spark.SparkFunSuite
import org.apache.spark.network.client.{TransportClient, TransportResponseHandler}
import org.apache.spark.network.server.StreamManager
import org.apache.spark.rpc._

class NettyRpcHandlerSuite extends SparkFunSuite {

  val env = mock(classOf[NettyRpcEnv])
  val sm = mock(classOf[StreamManager])
  when(env.deserialize(any(classOf[TransportClient]), any(classOf[ByteBuffer]))(any()))
    .thenReturn(RequestMessage(RpcAddress("localhost", 12345), null, null))

  test("receive") {
    val dispatcher = mock(classOf[Dispatcher])
    val nettyRpcHandler = new NettyRpcHandler(dispatcher, env, sm)

    val channel = mock(classOf[Channel])
    val client = new TransportClient(channel, mock(classOf[TransportResponseHandler]))
    when(channel.remoteAddress()).thenReturn(new InetSocketAddress("localhost", 40000))
    nettyRpcHandler.channelActive(client)

    verify(dispatcher, times(1)).postToAll(RemoteProcessConnected(RpcAddress("localhost", 40000)))
  }

  test("connectionTerminated") {
    val dispatcher = mock(classOf[Dispatcher])
    val nettyRpcHandler = new NettyRpcHandler(dispatcher, env, sm)

    val channel = mock(classOf[Channel])
    val client = new TransportClient(channel, mock(classOf[TransportResponseHandler]))
    when(channel.remoteAddress()).thenReturn(new InetSocketAddress("localhost", 40000))
    nettyRpcHandler.channelActive(client)

    when(channel.remoteAddress()).thenReturn(new InetSocketAddress("localhost", 40000))
    nettyRpcHandler.channelInactive(client)

    verify(dispatcher, times(1)).postToAll(RemoteProcessConnected(RpcAddress("localhost", 40000)))
    verify(dispatcher, times(1)).postToAll(
      RemoteProcessDisconnected(RpcAddress("localhost", 40000)))
  }

} 
Example 35
Source File: CustomRecoveryModeFactory.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
// This file is placed in different package to make sure all of these components work well
// when they are outside of org.apache.spark.
package other.supplier

import java.nio.ByteBuffer

import scala.collection.mutable
import scala.reflect.ClassTag

import org.apache.spark.SparkConf
import org.apache.spark.deploy.master._
import org.apache.spark.serializer.Serializer

class CustomRecoveryModeFactory(
  conf: SparkConf,
  serializer: Serializer
) extends StandaloneRecoveryModeFactory(conf, serializer) {

  CustomRecoveryModeFactory.instantiationAttempts += 1

  
  override def read[T: ClassTag](prefix: String): Seq[T] = {
    CustomPersistenceEngine.readAttempts += 1
    val results = for ((name, bytes) <- data; if name.startsWith(prefix))
      yield serializer.newInstance().deserialize[T](ByteBuffer.wrap(bytes))
    results.toSeq
  }
}

object CustomPersistenceEngine {
  @volatile var persistAttempts = 0
  @volatile var unpersistAttempts = 0
  @volatile var readAttempts = 0

  @volatile var lastInstance: Option[CustomPersistenceEngine] = None
}

class CustomLeaderElectionAgent(val masterInstance: LeaderElectable) extends LeaderElectionAgent {
  masterInstance.electedLeader()
} 
Example 36
Source File: GenericAvroSerializerSuite.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.serializer

import java.io.{ByteArrayInputStream, ByteArrayOutputStream}
import java.nio.ByteBuffer

import com.esotericsoftware.kryo.io.{Input, Output}
import org.apache.avro.{Schema, SchemaBuilder}
import org.apache.avro.generic.GenericData.Record

import org.apache.spark.{SharedSparkContext, SparkFunSuite}

class GenericAvroSerializerSuite extends SparkFunSuite with SharedSparkContext {
  conf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer")

  val schema : Schema = SchemaBuilder
    .record("testRecord").fields()
    .requiredString("data")
    .endRecord()
  val record = new Record(schema)
  record.put("data", "test data")

  test("schema compression and decompression") {
    val genericSer = new GenericAvroSerializer(conf.getAvroSchema)
    assert(schema === genericSer.decompress(ByteBuffer.wrap(genericSer.compress(schema))))
  }

  test("record serialization and deserialization") {
    val genericSer = new GenericAvroSerializer(conf.getAvroSchema)

    val outputStream = new ByteArrayOutputStream()
    val output = new Output(outputStream)
    genericSer.serializeDatum(record, output)
    output.flush()
    output.close()

    val input = new Input(new ByteArrayInputStream(outputStream.toByteArray))
    assert(genericSer.deserializeDatum(input) === record)
  }

  test("uses schema fingerprint to decrease message size") {
    val genericSerFull = new GenericAvroSerializer(conf.getAvroSchema)

    val output = new Output(new ByteArrayOutputStream())

    val beginningNormalPosition = output.total()
    genericSerFull.serializeDatum(record, output)
    output.flush()
    val normalLength = output.total - beginningNormalPosition

    conf.registerAvroSchemas(schema)
    val genericSerFinger = new GenericAvroSerializer(conf.getAvroSchema)
    val beginningFingerprintPosition = output.total()
    genericSerFinger.serializeDatum(record, output)
    val fingerprintLength = output.total - beginningFingerprintPosition

    assert(fingerprintLength < normalLength)
  }

  test("caches previously seen schemas") {
    val genericSer = new GenericAvroSerializer(conf.getAvroSchema)
    val compressedSchema = genericSer.compress(schema)
    val decompressedSchema = genericSer.decompress(ByteBuffer.wrap(compressedSchema))

    assert(compressedSchema.eq(genericSer.compress(schema)))
    assert(decompressedSchema.eq(genericSer.decompress(ByteBuffer.wrap(compressedSchema))))
  }
} 
Example 37
Source File: DiskStoreSuite.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.storage

import java.nio.{ByteBuffer, MappedByteBuffer}
import java.util.Arrays

import org.apache.spark.{SparkConf, SparkFunSuite}
import org.apache.spark.util.io.ChunkedByteBuffer
import org.apache.spark.util.Utils

class DiskStoreSuite extends SparkFunSuite {

  test("reads of memory-mapped and non memory-mapped files are equivalent") {
    // It will cause error when we tried to re-open the filestore and the
    // memory-mapped byte buffer tot he file has not been GC on Windows.
    assume(!Utils.isWindows)
    val confKey = "spark.storage.memoryMapThreshold"

    // Create a non-trivial (not all zeros) byte array
    val bytes = Array.tabulate[Byte](1000)(_.toByte)
    val byteBuffer = new ChunkedByteBuffer(ByteBuffer.wrap(bytes))

    val blockId = BlockId("rdd_1_2")
    val diskBlockManager = new DiskBlockManager(new SparkConf(), deleteFilesOnStop = true)

    val diskStoreMapped = new DiskStore(new SparkConf().set(confKey, "0"), diskBlockManager)
    diskStoreMapped.putBytes(blockId, byteBuffer)
    val mapped = diskStoreMapped.getBytes(blockId)
    assert(diskStoreMapped.remove(blockId))

    val diskStoreNotMapped = new DiskStore(new SparkConf().set(confKey, "1m"), diskBlockManager)
    diskStoreNotMapped.putBytes(blockId, byteBuffer)
    val notMapped = diskStoreNotMapped.getBytes(blockId)

    // Not possible to do isInstanceOf due to visibility of HeapByteBuffer
    assert(notMapped.getChunks().forall(_.getClass.getName.endsWith("HeapByteBuffer")),
      "Expected HeapByteBuffer for un-mapped read")
    assert(mapped.getChunks().forall(_.isInstanceOf[MappedByteBuffer]),
      "Expected MappedByteBuffer for mapped read")

    def arrayFromByteBuffer(in: ByteBuffer): Array[Byte] = {
      val array = new Array[Byte](in.remaining())
      in.get(array)
      array
    }

    assert(Arrays.equals(mapped.toArray, bytes))
    assert(Arrays.equals(notMapped.toArray, bytes))
  }
} 
Example 38
Source File: ChunkedByteBufferSuite.scala    From drizzle-spark   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.io

import java.nio.ByteBuffer

import com.google.common.io.ByteStreams

import org.apache.spark.SparkFunSuite
import org.apache.spark.network.util.ByteArrayWritableChannel
import org.apache.spark.util.io.ChunkedByteBuffer

class ChunkedByteBufferSuite extends SparkFunSuite {

  test("no chunks") {
    val emptyChunkedByteBuffer = new ChunkedByteBuffer(Array.empty[ByteBuffer])
    assert(emptyChunkedByteBuffer.size === 0)
    assert(emptyChunkedByteBuffer.getChunks().isEmpty)
    assert(emptyChunkedByteBuffer.toArray === Array.empty)
    assert(emptyChunkedByteBuffer.toByteBuffer.capacity() === 0)
    assert(emptyChunkedByteBuffer.toNetty.capacity() === 0)
    emptyChunkedByteBuffer.toInputStream(dispose = false).close()
    emptyChunkedByteBuffer.toInputStream(dispose = true).close()
  }

  test("getChunks() duplicates chunks") {
    val chunkedByteBuffer = new ChunkedByteBuffer(Array(ByteBuffer.allocate(8)))
    chunkedByteBuffer.getChunks().head.position(4)
    assert(chunkedByteBuffer.getChunks().head.position() === 0)
  }

  test("copy() does not affect original buffer's position") {
    val chunkedByteBuffer = new ChunkedByteBuffer(Array(ByteBuffer.allocate(8)))
    chunkedByteBuffer.copy(ByteBuffer.allocate)
    assert(chunkedByteBuffer.getChunks().head.position() === 0)
  }

  test("writeFully() does not affect original buffer's position") {
    val chunkedByteBuffer = new ChunkedByteBuffer(Array(ByteBuffer.allocate(8)))
    chunkedByteBuffer.writeFully(new ByteArrayWritableChannel(chunkedByteBuffer.size.toInt))
    assert(chunkedByteBuffer.getChunks().head.position() === 0)
  }

  test("toArray()") {
    val empty = ByteBuffer.wrap(Array.empty[Byte])
    val bytes = ByteBuffer.wrap(Array.tabulate(8)(_.toByte))
    val chunkedByteBuffer = new ChunkedByteBuffer(Array(bytes, bytes, empty))
    assert(chunkedByteBuffer.toArray === bytes.array() ++ bytes.array())
  }

  test("toArray() throws UnsupportedOperationException if size exceeds 2GB") {
    val fourMegabyteBuffer = ByteBuffer.allocate(1024 * 1024 * 4)
    fourMegabyteBuffer.limit(fourMegabyteBuffer.capacity())
    val chunkedByteBuffer = new ChunkedByteBuffer(Array.fill(1024)(fourMegabyteBuffer))
    assert(chunkedByteBuffer.size === (1024L * 1024L * 1024L * 4L))
    intercept[UnsupportedOperationException] {
      chunkedByteBuffer.toArray
    }
  }

  test("toInputStream()") {
    val empty = ByteBuffer.wrap(Array.empty[Byte])
    val bytes1 = ByteBuffer.wrap(Array.tabulate(256)(_.toByte))
    val bytes2 = ByteBuffer.wrap(Array.tabulate(128)(_.toByte))
    val chunkedByteBuffer = new ChunkedByteBuffer(Array(empty, bytes1, bytes2))
    assert(chunkedByteBuffer.size === bytes1.limit() + bytes2.limit())

    val inputStream = chunkedByteBuffer.toInputStream(dispose = false)
    val bytesFromStream = new Array[Byte](chunkedByteBuffer.size.toInt)
    ByteStreams.readFully(inputStream, bytesFromStream)
    assert(bytesFromStream === bytes1.array() ++ bytes2.array())
    assert(chunkedByteBuffer.getChunks().head.position() === 0)
  }
} 
Example 39
Source File: Main.scala    From ros_hadoop   with Apache License 2.0 5 votes vote down vote up
package de.valtech.foss

import scala.io.Source
import scala.collection.mutable.Map
import scala.collection.mutable.ListBuffer
import scala.collection.JavaConverters._
import Console.{GREEN, RED, RESET}
import scala.language.reflectiveCalls

import java.io.File
import java.io.FileInputStream
import java.io.FileOutputStream
import java.nio.channels.FileChannel.MapMode._
import java.nio.ByteOrder._
import java.nio.ByteBuffer

import de.valtech.foss.proto.RosbagIdxOuterClass.RosbagIdx

object Main extends App {
  def help() = {
    Console.err.printf(s"""
${RESET}${GREEN}Usage:
	--file <ros.bag> file to process
	--version print Rosbag version and exit
	--offset <offset> --number <records> Seek at offset < 1073741824 and read the specified number of records
${RESET}By default will just create the protobuf idx file needed for configuration.\n\n""")
    sys.exit(0)
  }

  val pargs = Map[String,AnyRef]()
  def process_cli(args: List[String]) :Boolean = args match {
    case Nil => true // parse success
    case "-v" :: rest => pargs += ("version" -> Some(true)); process_cli(rest)
    case "--version" :: rest => pargs += ("version" -> Some(true)); process_cli(rest)
    case "-f" :: x :: rest => pargs += ("file" -> x); process_cli(rest)
    case "--file" :: x :: rest => pargs += ("file" -> x); process_cli(rest)
    case "-n" :: x :: rest => pargs += ("number" -> Some(x.toInt)); process_cli(rest)
    case "--number" :: x :: rest => pargs += ("number" -> Some(x.toInt)); process_cli(rest)
    case "-o" :: x :: rest => pargs += ("offset" -> Some(x.toInt)); process_cli(rest)
    case "--offset" :: x :: rest => pargs += ("offset" -> Some(x.toInt)); process_cli(rest)
    case "-h" :: rest => help(); false
    case "--help" :: rest => help(); false
    case _ => Console.err.printf(s"${RESET}${RED}Unknown argument " + args.head); false
  }
  process_cli(args.toList)

  def use[T <: { def close() }]
    (resource: T)
    (code: T ⇒ Unit) =
    try
      code(resource)
    finally
      resource.close()

  pargs("file") match {
    case f:String => process()
    case _ => help()
  }

  def process(): Unit = {
    val fin = new File(pargs("file").asInstanceOf[String])
    use(new FileInputStream(fin)) { stream => {
      //printf("min: %s\n", Math.min(1073741824, fin.length) )
      val buffer = stream.getChannel.map(READ_ONLY, 0, Math.min(1073741824, fin.length)).order(LITTLE_ENDIAN)
      val p:RosbagParser = new RosbagParser(buffer)
      val version = p.read_version()
      val h = p.read_record().get
      if(pargs contains "version") {
        printf("%s\n%s\n\n", version, h)
        return
      }
      if(pargs contains "number"){
        buffer position pargs.getOrElse("offset",None).asInstanceOf[Option[Int]].getOrElse(0)
        for(i <- List.range(0,pargs("number").asInstanceOf[Option[Int]].getOrElse(0)))
          println(p.read_record)
        return
      }
      val idxpos = h.header.fields("index_pos").asInstanceOf[Long]
      //printf("idxpos: %s %s\n", idxpos, Math.min(1073741824, fin.length) )
      val b = stream.getChannel.map(READ_ONLY, idxpos, Math.min(1073741824, fin.length - idxpos)).order(LITTLE_ENDIAN)
      val pp:RosbagParser = new RosbagParser(b)
      val c = pp.read_connections(h.header, Nil)
      val chunk_idx = pp.read_chunk_infos(c)
      Console.err.printf(s"""${RESET}${GREEN}Found: """
          + chunk_idx.size
          +s""" chunks\n${RESET}It should be the same number reported by rosbag tool.\nIf you encounter any issues try reindexing your file and submit an issue.
          ${RESET}\n""")
      val fout = new FileOutputStream(pargs("file").asInstanceOf[String] + ".idx.bin")
      val builder = RosbagIdx.newBuilder
      for(i <- chunk_idx) builder.addArray(i)
      builder.build().writeTo(fout)
      fout.close()
      //printf("[%s]\n",chunk_idx.toArray.mkString(","))
    }}
  }
} 
Example 40
Source File: BGRImgToLocalSeqFile.scala    From BigDL   with Apache License 2.0 5 votes vote down vote up
package com.intel.analytics.bigdl.dataset.image

import java.nio.ByteBuffer
import java.nio.file.Path

import com.intel.analytics.bigdl.dataset.Transformer
import org.apache.hadoop.conf.Configuration
import org.apache.hadoop.fs.{Path => hadoopPath}
import org.apache.hadoop.io.{SequenceFile, Text}

import scala.collection.Iterator

object BGRImgToLocalSeqFile {
  def apply(blockSize: Int, baseFileName: Path, hasName: Boolean = false): BGRImgToLocalSeqFile = {
    new BGRImgToLocalSeqFile(blockSize, baseFileName, hasName)
  }
}


class BGRImgToLocalSeqFile(blockSize: Int, baseFileName: Path, hasName: Boolean = false) extends
  Transformer[(LabeledBGRImage, String), String] {
  private val conf: Configuration = new Configuration
  private var index = 0
  private val preBuffer: ByteBuffer = ByteBuffer.allocate(4 * 2)

  override def apply(prev: Iterator[(LabeledBGRImage, String)]): Iterator[String] = {
    new Iterator[String] {
      override def hasNext: Boolean = prev.hasNext

      override def next(): String = {
        val fileName = baseFileName + s"_$index.seq"
        val path = new hadoopPath(fileName)
        val writer = SequenceFile.createWriter(conf, SequenceFile.Writer.file(path),
          SequenceFile.Writer.keyClass(classOf[Text]),
          SequenceFile.Writer.valueClass(classOf[Text]))
        var i = 0
        while (i < blockSize && prev.hasNext) {
          val (image, imageName) = prev.next()

          preBuffer.putInt(image.width())
          preBuffer.putInt(image.height())
          val imageByteData = image.convertToByte()
          val data: Array[Byte] = new Array[Byte](preBuffer.capacity + imageByteData.length)
          System.arraycopy(preBuffer.array, 0, data, 0, preBuffer.capacity)
          System.arraycopy(imageByteData, 0, data, preBuffer.capacity, imageByteData.length)
          preBuffer.clear
          val imageKey = if (hasName) s"${imageName}\n${image.label().toInt}"
            else s"${image.label().toInt}"
          writer.append(new Text(imageKey), new Text(data))
          i += 1
        }
        writer.close()
        index += 1
        fileName
      }
    }
  }
} 
Example 41
Source File: BytesToBGRImg.scala    From BigDL   with Apache License 2.0 5 votes vote down vote up
package com.intel.analytics.bigdl.dataset.image

import java.awt.Color
import java.awt.image.{BufferedImage, DataBufferByte}
import java.nio.ByteBuffer

import com.intel.analytics.bigdl.dataset.{ByteRecord, Transformer}

import scala.collection.Iterator

object BytesToBGRImg {
  def apply(normalize: Float = 255f, resizeW : Int = -1, resizeH : Int = -1): BytesToBGRImg =
    new BytesToBGRImg(normalize, resizeW, resizeH)
}


class BytesToBGRImg(normalize: Float, resizeW : Int = -1, resizeH : Int = -1)
  extends Transformer[ByteRecord, LabeledBGRImage] {

  private val buffer = new LabeledBGRImage()

  override def apply(prev: Iterator[ByteRecord]): Iterator[LabeledBGRImage] = {
    prev.map(rawData => {
      buffer.copy(getImgData(rawData, resizeW, resizeH), normalize).setLabel(rawData.label)
    })
  }

  private def getImgData (record : ByteRecord, resizeW : Int, resizeH : Int)
  : Array[Byte] = {
    if (resizeW == -1) {
      return record.data
    } else {
      val rawData = record.data
      val imgBuffer = ByteBuffer.wrap(rawData)
      val width = imgBuffer.getInt
      val height = imgBuffer.getInt
      val bufferedImage : BufferedImage
      = new BufferedImage(width, height, BufferedImage.TYPE_3BYTE_BGR)
      val outputImagePixelData = bufferedImage.getRaster.getDataBuffer
        .asInstanceOf[DataBufferByte].getData
      System.arraycopy(imgBuffer.array(), 8,
        outputImagePixelData, 0, outputImagePixelData.length)
      BGRImage.resizeImage(bufferedImage, resizeW, resizeH)
    }
  }
} 
Example 42
Source File: TFRecordWriter.scala    From BigDL   with Apache License 2.0 5 votes vote down vote up
package com.intel.analytics.bigdl.utils.tf

import java.io.OutputStream
import java.nio.{ByteBuffer, ByteOrder}

import com.intel.analytics.bigdl.utils.Crc32

class TFRecordWriter(out: OutputStream) {

  private def toByteArrayAsLong(data: Long): Array[Byte] = {
    val buff = new Array[Byte](8)
    val bb = ByteBuffer.wrap(buff)
    bb.order(ByteOrder.LITTLE_ENDIAN)
    bb.putLong(data)
    buff
  }

  private def toByteArrayAsInt(data: Int): Array[Byte] = {
    val buff = new Array[Byte](4)
    val bb = ByteBuffer.wrap(buff)
    bb.order(ByteOrder.LITTLE_ENDIAN)
    bb.putInt(data)
    buff
  }

  def write(record: Array[Byte], offset: Int, length: Int): Unit = {
    val len = toByteArrayAsLong(length)
    out.write(len)
    out.write(toByteArrayAsInt(Crc32.maskedCRC32(len).toInt))
    out.write(record, offset, length)
    out.write(toByteArrayAsInt(Crc32.maskedCRC32(record, offset, length).toInt))
  }

  def write(record: Array[Byte]): Unit = {
    write(record, 0, record.length)
  }
} 
Example 43
Source File: TFRecordIterator.scala    From BigDL   with Apache License 2.0 5 votes vote down vote up
package com.intel.analytics.bigdl.utils.tf

import java.io.{BufferedInputStream, File, FileInputStream, InputStream}
import java.nio.{ByteBuffer, ByteOrder}


class TFRecordIterator(inputStream: InputStream) extends Iterator[Array[Byte]] {

  private var dataBuffer: Array[Byte] = null

  private val lengthBuffer: Array[Byte] = new Array[Byte](8)



  override def hasNext: Boolean = {
    if (dataBuffer != null) {
      true
    } else {
      val numOfBytes = inputStream.read(lengthBuffer)
      if (numOfBytes == 8) {
        val lengthWrapper = ByteBuffer.wrap(lengthBuffer)
        lengthWrapper.order(ByteOrder.LITTLE_ENDIAN)
        val length = lengthWrapper.getLong().toInt
        // todo, do crc check, simply skip now
        inputStream.skip(4)

        dataBuffer = new Array[Byte](length)
        inputStream.read(dataBuffer)
        // todo, do crc check, simply skip now
        inputStream.skip(4)
        true
      } else {
        inputStream.close()
        false
      }
    }
  }

  override def next(): Array[Byte] = {
    if (hasNext) {
      val data = this.dataBuffer
      this.dataBuffer = null
      data
    } else {
      throw new NoSuchElementException("next on empty iterator")
    }
  }
}

object TFRecordIterator {
  def apply(file: File): TFRecordIterator = {
    val inputStream = new FileInputStream(file)
    new TFRecordIterator(inputStream)
  }
} 
Example 44
Source File: Utils.scala    From BigDL   with Apache License 2.0 5 votes vote down vote up
package com.intel.analytics.bigdl.models.autoencoder

import java.nio.ByteBuffer
import java.nio.file.{Files, Path}

import com.intel.analytics.bigdl.dataset.ByteRecord

import scopt.OptionParser

object Utils {
  val trainMean = 0.13066047740239436
  val trainStd = 0.30810779333114624

  case class TrainParams(
    folder: String = "./",
    checkpoint: Option[String] = None,
    modelSnapshot: Option[String] = None,
    stateSnapshot: Option[String] = None,
    batchSize: Int = 150,
    maxEpoch: Int = 10,
    graphModel: Boolean = false,
    optimizerVersion: Option[String] = None
  )

  val trainParser = new OptionParser[TrainParams]("BigDL Autoencoder on MNIST") {
    opt[String]('f', "folder")
      .text("where you put the MNIST data")
      .action((x, c) => c.copy(folder = x))
    opt[String]("model")
      .text("model snapshot location")
      .action((x, c) => c.copy(modelSnapshot = Some(x)))
    opt[String]("state")
      .text("state snapshot location")
      .action((x, c) => c.copy(stateSnapshot = Some(x)))
    opt[String]("checkpoint")
      .text("where to cache the model and state")
      .action((x, c) => c.copy(checkpoint = Some(x)))
    opt[Int]('b', "batchSize")
      .text("batch size")
      .action((x, c) => c.copy(batchSize = x))
    opt[Int]('e', "maxEpoch")
      .text("max epoch")
      .action((x, c) => c.copy(maxEpoch = x))
    opt[Unit]('g', "graphModel")
      .text("use graph model")
      .action((x, c) => c.copy(graphModel = true))
    opt[String]("optimizerVersion")
      .text("state optimizer version")
      .action((x, c) => c.copy(optimizerVersion = Some(x)))
  }

  private[bigdl] def load(featureFile: Path, labelFile: Path): Array[ByteRecord] = {
    val labelBuffer = ByteBuffer.wrap(Files.readAllBytes(labelFile))
    val featureBuffer = ByteBuffer.wrap(Files.readAllBytes(featureFile))
    val labelMagicNumber = labelBuffer.getInt()

    require(labelMagicNumber == 2049)
    val featureMagicNumber = featureBuffer.getInt()
    require(featureMagicNumber == 2051)

    val labelCount = labelBuffer.getInt()
    val featureCount = featureBuffer.getInt()
    require(labelCount == featureCount)

    val rowNum = featureBuffer.getInt()
    val colNum = featureBuffer.getInt()

    val result = new Array[ByteRecord](featureCount)
    var i = 0
    while (i < featureCount) {
      val img = new Array[Byte]((rowNum * colNum))
      var y = 0
      while (y < rowNum) {
        var x = 0
        while (x < colNum) {
          img(x + y * colNum) = featureBuffer.get()
          x += 1
        }
        y += 1
      }
      result(i) = ByteRecord(img, labelBuffer.get().toFloat + 1.0f)
      i += 1
    }

    result
  }
} 
Example 45
Source File: FileReader.scala    From BigDL   with Apache License 2.0 5 votes vote down vote up
package com.intel.analytics.bigdl.visualization.tensorboard

import java.io.{BufferedInputStream}
import java.nio.ByteBuffer

import org.apache.hadoop.conf.Configuration
import org.apache.hadoop.fs.{FileSystem, Path}
import org.tensorflow.util.Event

import scala.collection.mutable.ArrayBuffer
import scala.util.matching.Regex

private[bigdl] object FileReader {
  val fileNameRegex = """bigdl.tfevents.*""".r

  
  def readScalar(file: Path, tag: String, fs: FileSystem): Array[(Long, Float, Double)] = {
    require(fs.isFile(file), s"FileReader: ${file} should be a file")
    val bis = new BufferedInputStream(fs.open(file))
    val longBuffer = new Array[Byte](8)
    val crcBuffer = new Array[Byte](4)
    val bf = new ArrayBuffer[(Long, Float, Double)]
    while (bis.read(longBuffer) > 0) {
      val l = ByteBuffer.wrap(longBuffer.reverse).getLong()
      bis.read(crcBuffer)
      // TODO: checksum
      //      val crc1 = ByteBuffer.wrap(crcBuffer.reverse).getInt()
      val eventBuffer = new Array[Byte](l.toInt)
      bis.read(eventBuffer)
      val e = Event.parseFrom(eventBuffer)
      if (e.getSummary.getValueCount == 1 &&
        tag.equals(e.getSummary.getValue(0).getTag())) {
        bf.append((e.getStep, e.getSummary.getValue(0).getSimpleValue,
          e.getWallTime))
      }
      bis.read(crcBuffer)
      //      val crc2 = ByteBuffer.wrap(crcBuffer.reverse).getInt()
    }
    bis.close()
    bf.toArray.sortWith(_._1 < _._1)
  }
} 
Example 46
Source File: BlockManagerWrapper.scala    From BigDL   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.storage

import java.nio.ByteBuffer

import org.apache.spark.SparkEnv

object BlockManagerWrapper {

  def putBytes( blockId: BlockId,
    bytes: ByteBuffer,
    level: StorageLevel): Unit = {
    require(bytes != null, "Bytes is null")
    SparkEnv.get.blockManager.putBytes(blockId, bytes, level)
  }

  def putSingle(blockId: BlockId,
    value: Any,
    level: StorageLevel,
    tellMaster: Boolean = true): Unit = {
    SparkEnv.get.blockManager.putSingle(blockId, value, level, tellMaster)
  }

  def removeBlock(blockId: BlockId): Unit = {
    SparkEnv.get.blockManager.removeBlock(blockId)
  }

  def getLocal(blockId: BlockId): Option[BlockResult] = {
    SparkEnv.get.blockManager.getLocal(blockId)
  }


  def getLocalBytes(blockId: BlockId): Option[ByteBuffer] = {
    SparkEnv.get.blockManager.getLocalBytes(blockId)
  }

  def getLocalOrRemoteBytes(blockId: BlockId): Option[ByteBuffer] = {
    val bm = SparkEnv.get.blockManager
    val maybeLocalBytes = bm.getLocalBytes(blockId)
    if (maybeLocalBytes.isDefined) {
      maybeLocalBytes
    } else {
      bm.getRemoteBytes(blockId)
    }
  }

  def unlock(blockId : BlockId): Unit = {}
} 
Example 47
Source File: ByteUtils.scala    From mantis   with Apache License 2.0 5 votes vote down vote up
package io.iohk.ethereum.utils

import java.math.BigInteger
import java.nio.{ByteBuffer, ByteOrder}

import akka.util.ByteString

import scala.util.Random

object ByteUtils {

  def bigIntegerToBytes(b: BigInteger, numBytes: Int): Array[Byte] = {
    val bytes = new Array[Byte](numBytes)
    val biBytes = b.toByteArray
    val start = if (biBytes.length == numBytes + 1) 1 else 0
    val length = Math.min(biBytes.length, numBytes)
    System.arraycopy(biBytes, start, bytes, numBytes - length, length)
    bytes
  }

  def xor(a: Array[Byte], b: Array[Byte]): Array[Byte] = {
    (a zip b) map { case (b1, b2) => (b1 ^ b2).toByte }
  }

  def or(arrays: Array[Byte]*): Array[Byte] = {
    require(arrays.map(_.length).distinct.length <= 1, "All the arrays should have the same length")
    require(arrays.nonEmpty, "There should be one or more arrays")

    val zeroes = Array.fill(arrays.head.length)(0.toByte)
    arrays.foldLeft[Array[Byte]](zeroes){
      case (prevOr, array) => prevOr.zip(array).map{ case (b1, b2) => (b1 | b2).toByte }
    }
  }

  def and(arrays: Array[Byte]*): Array[Byte] = {
    require(arrays.map(_.length).distinct.length <= 1, "All the arrays should have the same length")
    require(arrays.nonEmpty, "There should be one or more arrays")

    val ones = Array.fill(arrays.head.length)(0xFF.toByte)
    arrays.foldLeft[Array[Byte]](ones){
      case (prevOr, array) => prevOr.zip(array).map{ case (b1, b2) => (b1 & b2).toByte }
    }
  }

  def randomBytes(len: Int): Array[Byte] = {
    val arr = new Array[Byte](len)
    new Random().nextBytes(arr)
    arr
  }

  def bigEndianToShort(bs: Array[Byte]): Short = {
    val n = bs(0) << 8
    (n | bs(1) & 0xFF).toShort
  }

  def padLeft(bytes: ByteString, length: Int, byte: Byte = 0): ByteString = {
    val l = math.max(0, length - bytes.length)
    val fill = Seq.fill[Byte](l)(byte)
    fill ++: bytes
  }

  def compactPickledBytes(buffer: ByteBuffer): ByteString = {
    val data = Array.ofDim[Byte](buffer.limit)
    buffer.rewind()
    buffer.get(data)
    ByteString(data)
  }


  def bytesToInts(bytes: Array[Byte]): Array[Int] =
    bytes.grouped(4).map(getIntFromWord).toArray

  def intsToBytes(input: Array[Int]): Array[Byte] = {
    input.flatMap { i =>
      Array(
        (i & 0xFF).toByte,
        ((i >> 8) & 0xFF).toByte,
        ((i >> 16) & 0xFF).toByte,
        ((i >> 24) & 0xFF).toByte)
    }
  }

  def getIntFromWord(arr: Array[Byte]): Int = {
    ByteBuffer.wrap(arr, 0, 4).order(ByteOrder.LITTLE_ENDIAN).getInt
  }

} 
Example 48
Source File: IodbDataSource.scala    From mantis   with Apache License 2.0 5 votes vote down vote up
package io.iohk.ethereum.db.dataSource

import java.io.File
import java.nio.ByteBuffer
import java.util.concurrent.atomic.AtomicLong

import io.iohk.iodb.{ByteArrayWrapper, LSMStore}

class IodbDataSource private (lSMStore: LSMStore, keySize: Int, path: String) extends DataSource {

  import IodbDataSource._

  override def get(namespace: Namespace, key: Key): Option[Value] = {
    require(namespace.length + key.length <= keySize, "Wrong key size in IODB get")
    val keyPadded = padToKeySize(namespace, key, keySize).toArray
    lSMStore.get(ByteArrayWrapper(keyPadded)).map(v => v.data.toIndexedSeq)
  }

  override def update(namespace: Namespace, toRemove: Seq[Key], toUpsert: Seq[(Key, Value)]): DataSource = {
    require(
      toRemove.forall{ keyToRemove => namespace.length + keyToRemove.length <= keySize } &&
        toUpsert.forall{ case (keyToUpSert, _) => namespace.length + keyToUpSert.length <= keySize },
      "Wrong key size in IODB update"
    )
    val toRemovePadded = toRemove.map(key => padToKeySize(namespace, key, keySize))
    val toUpsertPadded = toUpsert.map{case (key, value) => padToKeySize(namespace, key, keySize) -> value}
    lSMStore.update(
      ByteArrayWrapper(storageVersionGen()),
      toRemovePadded.map(key => ByteArrayWrapper(key.toArray)),
      asStorables(toUpsertPadded))
    new IodbDataSource(lSMStore, keySize, path)
  }

  override def clear: DataSource = {
    destroy()
    IodbDataSource(path, keySize)
  }

  override def close(): Unit = lSMStore.close()

  override def destroy(): Unit = {
    try {
      close()
    } finally {
      val directoryDeletionSuccess = deleteDirectory(new File(path))
      assert(directoryDeletionSuccess, "Iodb folder destruction failed")
    }
  }

  private def asStorables(keyValues: Seq[(Key, Value)]): Seq[(ByteArrayWrapper, ByteArrayWrapper)] =
    keyValues.map{ case (key, value) => ByteArrayWrapper(key.toArray) -> ByteArrayWrapper(value.toArray) }
}


object IodbDataSource {
  val KeySizeWithoutNamespace = 32
  val KeySize = 33

  private val updateCounter = new AtomicLong(System.currentTimeMillis())

  private def storageVersionGen(): Array[Byte] = {
    ByteBuffer.allocate(java.lang.Long.SIZE / java.lang.Byte.SIZE).putLong(updateCounter.incrementAndGet()).array()
  }

  
  def apply(path: String, keySize: Int): IodbDataSource = {
    val dir: File = new File(path)
    val dirSetupSuccess = (dir.exists() && dir.isDirectory) || dir.mkdirs()
    assert(dirSetupSuccess, "Iodb folder creation failed")

    val lSMStore: LSMStore = new LSMStore(dir = dir, keySize = keySize)
    new IodbDataSource(lSMStore, keySize, path)
  }

  private def deleteDirectory(dir: File): Boolean = {
    require(dir.exists() && dir.isDirectory, "Trying to delete a file thats not a folder")
    val files = dir.listFiles()
    val filesDeletionSuccess = files.forall { f =>
      val deleted = if (f.isDirectory) deleteDirectory(f) else f.delete()
      deleted
    }
    filesDeletionSuccess && dir.delete()
  }

  private def padToKeySize(namespace: IndexedSeq[Byte], key: IndexedSeq[Byte], keySize: Int): IndexedSeq[Byte] =
    namespace ++ key.padTo(keySize - namespace.size, 0.toByte)
} 
Example 49
Source File: MatfastSerializer.scala    From MatRel   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.sql.matfast.util

import java.math.BigDecimal
import java.nio.ByteBuffer
import java.util.{HashMap => JavaHashMap}

import scala.reflect.ClassTag

import com.esotericsoftware.kryo.{Kryo, Serializer}
import com.esotericsoftware.kryo.io.{Input, Output}
import com.twitter.chill.ResourcePool

import org.apache.spark.{SparkConf, SparkEnv}
import org.apache.spark.serializer.{KryoSerializer, SerializerInstance}
import org.apache.spark.sql.matfast.matrix._
import org.apache.spark.sql.types.Decimal
import org.apache.spark.util.MutablePair


private[matfast] class MatfastSerializer(conf: SparkConf) extends KryoSerializer(conf) {
  override def newKryo(): Kryo = {
    val kryo = super.newKryo()
    kryo.setRegistrationRequired(false)
    kryo.register(classOf[MutablePair[_, _]])
    kryo.register(classOf[org.apache.spark.sql.catalyst.expressions.GenericRow])
    kryo.register(classOf[org.apache.spark.sql.catalyst.expressions.GenericInternalRow])
    kryo.register(classOf[java.math.BigDecimal], new JavaBigDecimalSerializer)
    kryo.register(classOf[BigDecimal], new ScalaBigDecimalSerializer)

    kryo.register(classOf[Decimal])
    kryo.register(classOf[JavaHashMap[_, _]])
    kryo.register(classOf[DenseMatrix])
    kryo.register(classOf[SparseMatrix])

    kryo.setReferences(false)
    kryo
  }
}

private[matfast] class KryoResourcePool(size: Int) extends ResourcePool[SerializerInstance](size) {
  val ser: MatfastSerializer = {
    val sparkConf = Option(SparkEnv.get).map(_.conf).getOrElse(new SparkConf())
    new MatfastSerializer(sparkConf)
  }

  def newInstance(): SerializerInstance = ser.newInstance()
}

private[matfast] object MatfastSerializer {
  @transient lazy val resourcePool = new KryoResourcePool(50)

  private[this] def acquireRelease[O](fn: SerializerInstance => O): O = {
    val kryo = resourcePool.borrow()
    try {
      fn(kryo)
    } finally {
      resourcePool.release(kryo)
    }
  }

  def serialize[T: ClassTag](o: T): Array[Byte] = {
    acquireRelease { k =>
      k.serialize(o).array()
    }
  }

  def deserialize[T: ClassTag](bytes: Array[Byte]): T =
    acquireRelease { k =>
      k.deserialize[T](ByteBuffer.wrap(bytes))
    }
}

private[matfast] class JavaBigDecimalSerializer extends Serializer[java.math.BigDecimal] {
  def write(kryo: Kryo, output: Output, bd: java.math.BigDecimal) {
    output.writeString(bd.toString)
  }

  def read(kryo: Kryo, input: Input, tpe: Class[java.math.BigDecimal]): java.math.BigDecimal = {
    new java.math.BigDecimal(input.readString())
  }
}

private[matfast] class ScalaBigDecimalSerializer extends Serializer[BigDecimal] {
  def write(kryo: Kryo, output: Output, bd: BigDecimal): Unit = {
    output.writeString(bd.toString)
  }

  def read(kryo: Kryo, input: Input, tpe: Class[BigDecimal]): BigDecimal = {
    new java.math.BigDecimal(input.readString())
  }
} 
Example 50
Source File: RocksHelper.scala    From incubator-s2graph   with Apache License 2.0 5 votes vote down vote up
package org.apache.s2graph.core.storage.rocks

import java.nio.{ByteBuffer, ByteOrder}

import org.apache.s2graph.core.QueryParam

object RocksHelper {

  def intToBytes(value: Int): Array[Byte] = {
    val intBuffer = ByteBuffer.allocate(4).order(ByteOrder.nativeOrder())
    intBuffer.clear()
    intBuffer.putInt(value)
    intBuffer.array()
  }

  def longToBytes(value: Long): Array[Byte] = {
    val longBuffer = ByteBuffer.allocate(8).order(ByteOrder.nativeOrder())
    longBuffer.clear()
    longBuffer.putLong(value)
    longBuffer.array()
  }

  def bytesToInt(data: Array[Byte], offset: Int): Int = {
    if (data != null) {
      val intBuffer = ByteBuffer.allocate(4).order(ByteOrder.nativeOrder())
      intBuffer.put(data, offset, 4)
      intBuffer.flip()
      intBuffer.getInt()
    } else 0
  }

  def bytesToLong(data: Array[Byte], offset: Int): Long = {
    if (data != null) {
      val longBuffer = ByteBuffer.allocate(8).order(ByteOrder.nativeOrder())
      longBuffer.put(data, offset, 8)
      longBuffer.flip()
      longBuffer.getLong()
    } else 0L
  }

  case class ScanWithRange(cf: Array[Byte], startKey: Array[Byte], stopKey: Array[Byte], offset: Int, limit: Int)
  case class GetRequest(cf: Array[Byte], key: Array[Byte])

  type RocksRPC = Either[GetRequest, ScanWithRange]
} 
Example 51
Source File: OrderBookSideSnapshotCodecs.scala    From matcher   with MIT License 5 votes vote down vote up
package com.wavesplatform.dex.codecs

import java.math.BigInteger
import java.nio.ByteBuffer

import com.google.common.primitives.{Ints, Longs}
import com.wavesplatform.dex.codecs.ByteBufferCodecs.ByteBufferExt
import com.wavesplatform.dex.domain.model.Price
import com.wavesplatform.dex.domain.order.{Order, OrderType}
import com.wavesplatform.dex.model.{BuyLimitOrder, LimitOrder, OrderBookSideSnapshot, SellLimitOrder}

import scala.collection.mutable

object OrderBookSideSnapshotCodecs {

  def encode(dest: mutable.ArrayBuilder[Byte], snapshot: OrderBookSideSnapshot): Unit = {
    dest ++= Ints.toByteArray(snapshot.size)
    snapshot.foreach {
      case (price, xs) =>
        dest ++= Longs.toByteArray(price)
        dest ++= Ints.toByteArray(xs.size)
        xs.foreach(encodeLoV2(dest, _))
    }
  }

  def decode(bb: ByteBuffer): OrderBookSideSnapshot = {
    val snapshotSize = bb.getInt
    val r            = Map.newBuilder[Price, Seq[LimitOrder]]
    (1 to snapshotSize).foreach { _ =>
      val price       = bb.getLong
      val levelSize   = bb.getInt
      val limitOrders = (1 to levelSize).map(_ => decodeLo(bb))
      r += price -> limitOrders
    }
    r.result()
  }

  def encodeLoV1(dest: mutable.ArrayBuilder[Byte], lo: LimitOrder): Unit = {

    dest ++= lo.order.orderType.bytes
    dest ++= Longs.toByteArray(lo.amount)
    dest ++= Longs.toByteArray(lo.fee)
    dest += lo.order.version

    val orderBytes = lo.order.bytes()

    dest ++= Ints.toByteArray(orderBytes.length)
    dest ++= orderBytes
  }

  def encodeLoV2(dest: mutable.ArrayBuilder[Byte], lo: LimitOrder): Unit = {
    val avgWeighedPriceNominatorBytes = lo.avgWeighedPriceNominator.toByteArray

    dest += 2

    encodeLoV1(dest, lo)

    dest ++= Ints.toByteArray(avgWeighedPriceNominatorBytes.length)
    dest ++= avgWeighedPriceNominatorBytes
  }

  def decodeLo(bb: ByteBuffer): LimitOrder = {

    val header    = bb.get
    val version   = if (header == 2) 2 else 1
    val orderType = if (version == 1) header else bb.get

    val amount       = bb.getLong
    val fee          = bb.getLong
    val orderVersion = bb.get
    val order        = Order.fromBytes(orderVersion, bb.getBytes)

    val avgWeighedPriceNominator =
      if (version == 2) new BigInteger(bb.getBytes)
      else {
        val filledAmount = order.amount - amount
        (BigInt(order.price) * filledAmount).bigInteger
      }

    OrderType(orderType) match {
      case OrderType.SELL => SellLimitOrder(amount, fee, order, avgWeighedPriceNominator)
      case OrderType.BUY  => BuyLimitOrder(amount, fee, order, avgWeighedPriceNominator)
    }
  }
} 
Example 52
Source File: ByteBufferCodecs.scala    From matcher   with MIT License 5 votes vote down vote up
package com.wavesplatform.dex.codecs

import java.nio.ByteBuffer

import com.wavesplatform.dex.domain.asset.Asset
import com.wavesplatform.dex.domain.asset.Asset.{IssuedAsset, Waves}
import com.wavesplatform.dex.domain.bytes.ByteStr
import com.wavesplatform.dex.model.{AcceptedOrderType, OrderStatus}

object ByteBufferCodecs {

  implicit class ByteBufferExt(val b: ByteBuffer) extends AnyVal {

    def getBytes: Array[Byte] = {
      val len = b.getInt
      if (b.limit() < len || len < 0) {
        throw new Exception(s"Invalid array size ($len)")
      }
      val bytes = new Array[Byte](len)
      b.get(bytes)
      bytes
    }

    def putAssetId(assetId: Asset): ByteBuffer = assetId match {
      case Waves => b.put(0.toByte)
      case IssuedAsset(aid) =>
        require(aid.arr.length < Byte.MaxValue, "Asset ID is too long")
        b.put(aid.arr.length.toByte).put(aid.arr)
    }

    def getAssetId: Asset = b.get() match {
      case 0 => Waves
      case len =>
        val arr = new Array[Byte](len)
        b.get(arr)
        IssuedAsset(ByteStr(arr))
    }

    def putFinalOrderStatus(orderInfoVersion: Byte, st: OrderStatus): ByteBuffer = {
      val tpe: Byte = st match {
        case _: OrderStatus.Filled    => 0
        case _: OrderStatus.Cancelled => 1
        case x                        => throw new IllegalArgumentException(s"Can't encode order status $x")
      }
      val r = b.put(tpe).putLong(st.filledAmount)
      if (orderInfoVersion <= 1) r else r.putLong(st.filledFee)
    }

    def getFinalOrderStatus(orderInfoVersion: Byte, totalAmount: Long, totalFee: Long): OrderStatus.Final = {
      def fee(filledAmount: Long) = if (orderInfoVersion <= 1) (BigInt(filledAmount) * totalFee / totalAmount).toLong else b.getLong

      b.get match {
        case 0 =>
          val filledAmount = b.getLong
          OrderStatus.Filled(filledAmount, fee(filledAmount))
        case 1 =>
          val filledAmount = b.getLong
          OrderStatus.Cancelled(filledAmount, fee(filledAmount))
        case x => throw new IllegalArgumentException(s"Can't parse order status: $x")
      }
    }

    def putAcceptedOrderType(x: AcceptedOrderType): ByteBuffer = x match {
      case AcceptedOrderType.Limit  => b.put(0: Byte)
      case AcceptedOrderType.Market => b.put(1: Byte)
    }

    def getAcceptedOrderType: AcceptedOrderType = b.get match {
      case 0 => AcceptedOrderType.Limit
      case 1 => AcceptedOrderType.Market
      case x => throw new IllegalArgumentException(s"Can't parse accepted order type: $x")
    }
  }
} 
Example 53
Source File: OrderBookSnapshot.scala    From matcher   with MIT License 5 votes vote down vote up
package com.wavesplatform.dex.model

import java.nio.ByteBuffer

import com.wavesplatform.dex.codecs.OrderBookSideSnapshotCodecs

import scala.collection.mutable

case class OrderBookSnapshot(bids: OrderBookSideSnapshot, asks: OrderBookSideSnapshot, lastTrade: Option[LastTrade])
object OrderBookSnapshot {

  val empty: OrderBookSnapshot = OrderBookSnapshot(bids = Map.empty, asks = Map.empty, None)

  def serialize(dest: mutable.ArrayBuilder[Byte], x: OrderBookSnapshot): Unit = {
    OrderBookSideSnapshotCodecs.encode(dest, x.bids)
    OrderBookSideSnapshotCodecs.encode(dest, x.asks)
    x.lastTrade match {
      case None => dest += 0
      case Some(lastTrade) =>
        dest += 1
        LastTrade.serialize(dest, lastTrade)
    }
  }

  def fromBytes(bb: ByteBuffer): OrderBookSnapshot =
    OrderBookSnapshot(
      OrderBookSideSnapshotCodecs.decode(bb),
      OrderBookSideSnapshotCodecs.decode(bb),
      bb.get match {
        case 0 => None
        case 1 => Some(LastTrade.fromBytes(bb))
        case x => throw new RuntimeException(s"Can't deserialize Option as $x")
      }
    )
} 
Example 54
Source File: ClientHandler.scala    From spark-riak-connector   with Apache License 2.0 5 votes vote down vote up
package com.basho.riak.stub

import java.nio.ByteBuffer
import java.nio.channels.{AsynchronousSocketChannel, Channel, CompletionHandler}

import com.basho.riak.client.core.netty.RiakMessageCodec
import com.basho.riak.stub.ClientHandler._
import org.slf4j.LoggerFactory

class ClientHandler(val messageHandler: RiakMessageHandler) extends RiakMessageCodec
  with CompletionHandler[Integer, (AsynchronousSocketChannel, ByteBuffer)] {

  override def completed(result: Integer, attachment: (AsynchronousSocketChannel, ByteBuffer)): Unit = attachment match {
    case (channel, buffer) if result > 0 =>
      logger.info(s"Message received ${SocketUtils.serverConnectionAsStr(channel)} ($result bytes).")
      RiakMessageEncoder.decode(buffer.rewind().asInstanceOf[ByteBuffer]) match {
        case Some(m) if channel.isOpen =>
          val msgs = messageHandler.handle(new Context(channel), m)
          val encoded = RiakMessageEncoder.encode(msgs.toSeq: _*)
          val bytes = channel.write(encoded).get
          assert(bytes == encoded.position())
          logger.info(s"Response sent ${SocketUtils.clientConnectionAsStr(channel)} ($bytes bytes).")
          messageHandler.onRespond(m, msgs)
        case Some(m) if !channel.isOpen =>
          logger.warn("Impossible to write message to channel: channel has been already closed")
        case None => // TODO: handle case with no message
      }
      buffer.clear()
      channel.read(buffer, (channel, buffer), this)
    case _ =>
  }

  override def failed(exc: Throwable, attachment: (AsynchronousSocketChannel, ByteBuffer)): Unit = attachment match {
    case (channel, _) if channel.isOpen =>
      logger.error(s"Something went wrong with client ${SocketUtils.serverConnectionAsStr(channel)}", exc)
      disconnectClient(channel)
    case _ => // channel is already closed - do nothing
  }

  def disconnectClient(client: AsynchronousSocketChannel): Unit = this.synchronized {
    client.isOpen match {
      case true =>
        val connectionString = SocketUtils.serverConnectionAsStr(client)
        client.shutdownInput()
        client.shutdownOutput()
        client.close()
        logger.info(s"Client $connectionString was gracefully disconnected")
      case false => // client is already closed - do nothing
    }
  }
}

object ClientHandler {
  val logger = LoggerFactory.getLogger(classOf[ClientHandler])

  
  class Context(val channel: Channel)

} 
Example 55
Source File: ProxyMessageHandler.scala    From spark-riak-connector   with Apache License 2.0 5 votes vote down vote up
package com.basho.riak.stub

import java.net.InetSocketAddress
import java.nio.ByteBuffer
import java.nio.channels._

import com.basho.riak.client.core.RiakMessage
import com.basho.riak.client.core.util.HostAndPort
import shaded.com.basho.riak.protobuf.RiakKvPB
import shaded.com.basho.riak.protobuf.RiakMessageCodes._
import shaded.com.google.protobuf.ByteString

import scala.collection.JavaConversions._
import scala.collection.mutable.ArrayBuffer


class ProxyMessageHandler(hostAndPort: HostAndPort) extends RiakMessageHandler {

  private final val riakAddress = new InetSocketAddress(hostAndPort.getHost, hostAndPort.getPort)

  override def handle(context: ClientHandler.Context, input: RiakMessage): Iterable[RiakMessage] = input.getCode match {
    // coverage plan received from real Riak node must be modified to replace real node's host and port with proxy
    case MSG_CoverageReq => forwardAndTransform(context, input) { output =>
      val resp = RiakKvPB.RpbCoverageResp.parseFrom(output.getData)
      val modified = RiakKvPB.RpbCoverageResp.newBuilder(resp)
        .clearEntries()
        .addAllEntries(resp.getEntriesList.map { ce =>
          val ceBuilder = RiakKvPB.RpbCoverageEntry.newBuilder(ce)
          if (ce.getIp.toStringUtf8 == hostAndPort.getHost && ce.getPort == hostAndPort.getPort) {
            val localAddress = context.channel.asInstanceOf[NetworkChannel]
              .getLocalAddress.asInstanceOf[InetSocketAddress]
            ceBuilder.setIp(ByteString.copyFromUtf8(localAddress.getHostString))
            ceBuilder.setPort(localAddress.getPort)
          }
          ceBuilder.build()
        }).build()
      new RiakMessage(output.getCode, modified.toByteArray)
    }
    case _ => forwardMessage(context, input)
  }

  private def forwardMessage(context: ClientHandler.Context, input: RiakMessage): Iterable[RiakMessage] = {
    def readRiakResponse(channel: SocketChannel, out: List[RiakMessage] = Nil): Iterable[RiakMessage] = out match {
      case _ if !isDoneReceived(out, input) => readRiakResponse(channel, out ++ readSocket(channel))
      case _ => out
    }

    val channel = SocketChannel.open(riakAddress)
    try {
      // forward request to real Riak node
      assert(channel.write(RiakMessageEncoder.encode(input)) > 0)

      // read response for forwarded request from real Riak node
      readRiakResponse(channel)
    } finally {
      channel.close()
    }
  }

  private def readSocket(channel: SocketChannel): Iterable[RiakMessage] = {
    var accumulator = ByteBuffer.allocateDirect(0)

    var out = ArrayBuffer[RiakMessage]()
    while (out.isEmpty || accumulator.hasRemaining) {
      // try to parse riak message from bytes in accumulator buffer
      RiakMessageEncoder.decode(accumulator) match {
        case Some(x) =>
          accumulator = accumulator.slice()
          out += x
        case None =>
          // read next chunk of data from channel and add it into accumulator
          val in = ByteBuffer.allocateDirect(1024) // scalastyle:ignore
          channel.read(in)
          accumulator = ByteBuffer
            .allocate(accumulator.rewind().limit() + in.flip().limit())
            .put(accumulator)
            .put(in)
          accumulator.rewind()
          in.clear()
      }
    }
    out
  }

  private def isDoneReceived(out: Iterable[RiakMessage], input: RiakMessage): Boolean = input.getCode match {
    case MSG_IndexReq => out.foldLeft[Boolean](false)((a, m) => a || RiakKvPB.RpbIndexResp.parseFrom(m.getData).getDone)
    case _ => out.nonEmpty
  }

  private def forwardAndTransform(context: ClientHandler.Context, input: RiakMessage
                                 )(transform: RiakMessage => RiakMessage
                                 ): Iterable[RiakMessage] = forwardMessage(context, input).map(transform(_))

  override def onRespond(input: RiakMessage, output: Iterable[RiakMessage]): Unit = {}
} 
Example 56
Source File: RiakNodeStub.scala    From spark-riak-connector   with Apache License 2.0 5 votes vote down vote up
package com.basho.riak.stub

import java.net.InetSocketAddress
import java.nio.ByteBuffer
import java.nio.channels.{AsynchronousCloseException, AsynchronousServerSocketChannel, AsynchronousSocketChannel, CompletionHandler}

import com.basho.riak.client.core.util.HostAndPort
import com.basho.riak.stub.RiakNodeStub._
import org.slf4j.LoggerFactory

class RiakNodeStub(val host: String, val port: Int, messageHandler: RiakMessageHandler) {

  private final val localAddress = new InetSocketAddress(host, port)
  private final val clientHandler = new ClientHandler(messageHandler)

  private var serverChannel: AsynchronousServerSocketChannel = _
  private var clients: List[AsynchronousSocketChannel] = Nil

  def start(): HostAndPort = {
    serverChannel = AsynchronousServerSocketChannel.open()
    require(serverChannel.isOpen)

    serverChannel.bind(localAddress)
    serverChannel.accept(serverChannel, new CompletionHandler[AsynchronousSocketChannel, AsynchronousServerSocketChannel]() {
      override def completed(client: AsynchronousSocketChannel, server: AsynchronousServerSocketChannel): Unit = {
        logger.info(s"Incoming connection: ${SocketUtils.serverConnectionAsStr(client)}")
        this.synchronized {
          clients = client :: clients
        }

        val buffer = ByteBuffer.allocateDirect(1024) // scalastyle:ignore
        client.read(buffer, (client, buffer), clientHandler)

        server.accept(server, this)
      }

      override def failed(exc: Throwable, serverChannel: AsynchronousServerSocketChannel): Unit = exc match {
        case _: AsynchronousCloseException =>
        case _ => logger.error(s"Something went wrong:  ${serverChannel.toString}", exc);
      }
    })

    HostAndPort.fromParts(
      serverChannel.getLocalAddress.asInstanceOf[InetSocketAddress].getHostString,
      serverChannel.getLocalAddress.asInstanceOf[InetSocketAddress].getPort)
  }

  def stop(): Unit = this.synchronized {
    Option(serverChannel).foreach(_.close)
    clients.foreach(clientHandler.disconnectClient)
  }
}

object RiakNodeStub {
  val logger = LoggerFactory.getLogger(classOf[RiakNodeStub])
  final val DEFAULT_HOST = "localhost"

  def apply(host: String, port: Int, messageHandler: RiakMessageHandler): RiakNodeStub = new RiakNodeStub(host, port, messageHandler)

  def apply(port: Int, messageHandler: RiakMessageHandler): RiakNodeStub = RiakNodeStub(DEFAULT_HOST, port, messageHandler)

  def apply(messageHandler: RiakMessageHandler): RiakNodeStub = RiakNodeStub(DEFAULT_HOST, 0, messageHandler)
} 
Example 57
Source File: RiakMessageEncoder.scala    From spark-riak-connector   with Apache License 2.0 5 votes vote down vote up
package com.basho.riak.stub

import java.nio.ByteBuffer
import java.util

import com.basho.riak.client.core.RiakMessage
import com.basho.riak.client.core.netty.RiakMessageCodec
import io.netty.buffer.Unpooled

import scala.collection.JavaConversions._

object RiakMessageEncoder extends RiakMessageCodec {

  def encode(messages: RiakMessage*): ByteBuffer = {
    val buffer = Unpooled.buffer(messages.head.getData.length)
    messages.foreach(encode(null, _, buffer)) // scalastyle:ignore
    buffer.nioBuffer()
  }

  def decode(in: ByteBuffer): Option[RiakMessage] = {
    val out = new util.ArrayList[AnyRef]()
    val buffer = Unpooled.wrappedBuffer(in)
    decode(null, buffer, out) // scalastyle:ignore
    List(out: _*) match {
      case msg :: _ =>
        in.position(buffer.readerIndex())
        Some(msg.asInstanceOf[RiakMessage])
      case Nil => None
    }
  }
} 
Example 58
Source File: KvStreamingTest.scala    From spark-riak-connector   with Apache License 2.0 5 votes vote down vote up
package com.basho.riak.spark.streaming

import java.nio.ByteBuffer
import java.util.concurrent._

import com.basho.riak.spark._
import com.basho.riak.spark.rdd.{AbstractRiakSparkTest, RiakTSTests}
import org.junit.Assert.assertEquals
import org.junit.experimental.categories.Category
import org.junit.{After, Before, Test}

@Category(Array(classOf[RiakTSTests]))
class KvStreamingTest extends AbstractRiakSparkTest with SparkStreamingFixture {

  protected final val VALUES_NUMBER = 5

  protected final val executorService = Executors.newCachedThreadPool()
  private val dataSource = new SocketStreamingDataSource
  private var port = -1

  @Before
  def setUp(): Unit = {
    port = dataSource.start(client => {
      (0 until VALUES_NUMBER).foreach(i => client.write(ByteBuffer.wrap(s"value-$i\n".getBytes)))
      logInfo(s"$VALUES_NUMBER values were send to client")
    })
  }

  @After
  def tearDown(): Unit = {
    dataSource.stop()
  }

  @Test(timeout = 10 * 1000) // 10 seconds timeout
  def saveToRiak(): Unit = {
    executorService.submit(new Runnable {
      override def run(): Unit = {
        ssc.socketTextStream("localhost", port)
          .saveToRiak(DEFAULT_NAMESPACE_4STORE.getBucketNameAsString)

        ssc.start()
        ssc.awaitTerminationOrTimeout(5 * 1000)
      }
    })

    val data: List[(String, Any)] = executorService.submit(new Callable[List[(String, Any)]] {
      override def call(): List[(String, Any)] = {
        var rdd = sc.riakBucket(DEFAULT_NAMESPACE_4STORE).queryAll()
        var count = rdd.count()
        while (count < VALUES_NUMBER) {
          TimeUnit.SECONDS.sleep(2)

          rdd = sc.riakBucket(DEFAULT_NAMESPACE_4STORE).queryAll()
          count = rdd.count()
        }
        rdd.collect().toList
      }
    }).get()

    assertEquals(VALUES_NUMBER, data.length)
    assertEqualsUsingJSONIgnoreOrder(
      """[
        | ['${json-unit.ignore}', 'value-0'],
        | ['${json-unit.ignore}', 'value-1'],
        | ['${json-unit.ignore}', 'value-2'],
        | ['${json-unit.ignore}', 'value-3'],
        | ['${json-unit.ignore}', 'value-4']
        | ]""".stripMargin, data)
  }
} 
Example 59
Source File: Decimals.scala    From parquet4s   with MIT License 5 votes vote down vote up
package com.github.mjakubowski84.parquet4s

import java.math.MathContext
import java.nio.ByteBuffer

import org.apache.parquet.io.api.Binary

object Decimals {
  val Scale = 18
  val Precision = 38
  val ByteArrayLength = 16
  val MathContext = new MathContext(Precision)

  private def rescale(original: BigDecimal): BigDecimal = {
    if (original.scale == Scale && original.mc == MathContext) original
    else BigDecimal.decimal(original.bigDecimal, MathContext).setScale(Scale, BigDecimal.RoundingMode.HALF_UP)
  }

  def rescaleBinary(binary: Binary, originalScale: Int, originalMathContext: MathContext): Binary =
    binaryFromDecimal(decimalFromBinary(binary, originalScale, originalMathContext))

  def decimalFromBinary(binary: Binary, scale: Int = Scale, mathContext: MathContext = MathContext): BigDecimal =
    BigDecimal(BigInt(binary.getBytes), scale, mathContext)

  def binaryFromDecimal(decimal: BigDecimal): Binary = {
    
    val buf = ByteBuffer.allocate(ByteArrayLength)
    val unscaled = rescale(decimal).bigDecimal.unscaledValue().toByteArray
    // BigInteger is stored in tail of byte array, sign is stored in unoccupied cells
    val sign: Byte = if (unscaled.head < 0) -1 else 0
    (0 until ByteArrayLength - unscaled.length).foreach(_ => buf.put(sign))
    buf.put(unscaled)
    Binary.fromReusedByteArray(buf.array())
  }
} 
Example 60
Source File: ConversionUtils.scala    From scala-serialization   with MIT License 5 votes vote down vote up
package com.komanov.serialization.converters

import java.nio.ByteBuffer
import java.time.Instant
import java.util.UUID

import com.google.protobuf.ByteString

object ConversionUtils {

  def uuidToBytes(uuid: UUID): ByteString = {
    val bb = uuidToByteBuffer(uuid)
    if (bb == null) ByteString.EMPTY else ByteString.copyFrom(bb)
  }

  def uuidToByteBuffer(uuid: UUID): ByteBuffer = {
    if (uuid == null) {
      return null
    }

    val buffer = ByteBuffer.allocate(16)
    buffer.putLong(uuid.getMostSignificantBits)
    buffer.putLong(uuid.getLeastSignificantBits)
    buffer.rewind()
    buffer
  }

  def bytesToUuid(bb: ByteBuffer): UUID = {
    if (bb == null) {
      return null
    }

    val length = bb.limit() - bb.position()
    if (length == 0) {
      return null
    }

    require(length >= 16, s"expected 16 bytes: ${bb.capacity()} / ${bb.limit()}")

    new UUID(bb.getLong, bb.getLong)
  }

  def bytesToUuid(bs: ByteString): UUID = {
    bytesToUuid(bs.asReadOnlyByteBuffer())
  }

  def instantToLong(v: Instant) = v.toEpochMilli

  def longToInstance(v: Long) = Instant.ofEpochMilli(v)

} 
Example 61
Source File: BoopickleConverter.scala    From scala-serialization   with MIT License 5 votes vote down vote up
package com.komanov.serialization.converters

import java.nio.ByteBuffer
import java.time.Instant
import java.util

import boopickle.Default._
import boopickle.{BufferPool, DecoderSize, EncoderSize}
import com.komanov.serialization.domain._


object BoopickleConverter extends MyConverter {

  implicit def pickleState = new PickleState(new EncoderSize, false, false)
  implicit val unpickleState = (bb: ByteBuffer) => new UnpickleState(new DecoderSize(bb), false, false)

  BufferPool.disable()

  override def toByteArray(site: Site): Array[Byte] = {
    val bb = Pickle.intoBytes(site)
    val a = bbToArray(bb)
    BufferPool.release(bb)
    a
  }

  override def fromByteArray(bytes: Array[Byte]): Site = {
    Unpickle[Site].fromBytes(ByteBuffer.wrap(bytes))
  }

  override def toByteArray(event: SiteEvent): Array[Byte] = {
    val bb = Pickle.intoBytes(event)
    val a = bbToArray(bb)
    BufferPool.release(bb)
    a
  }

  override def siteEventFromByteArray(clazz: Class[_], bytes: Array[Byte]): SiteEvent = {
    Unpickle[SiteEvent].fromBytes(ByteBuffer.wrap(bytes))
  }

  private def bbToArray(bb: ByteBuffer) = {
    util.Arrays.copyOfRange(bb.array(), 0, bb.limit())
  }

  implicit val instantPickler = transformPickler[Instant, Long](t => Instant.ofEpochMilli(t))(_.toEpochMilli)

  implicit val pageComponentTypePickler = transformPickler(PageComponentType.valueOf)(_.name())
  implicit val siteFlagPickler = transformPickler(SiteFlag.valueOf)(_.name())
  implicit val siteTypePickler = transformPickler(SiteType.valueOf)(_.name())

  implicit val entryPointPickler = compositePickler[EntryPoint]
    .addConcreteType[DomainEntryPoint]
    .addConcreteType[FreeEntryPoint]

  implicit val pageComponentDataPickler = compositePickler[PageComponentData]
    .addConcreteType[TextComponentData]
    .addConcreteType[ButtonComponentData]
    .addConcreteType[BlogComponentData]

  implicit val siteEventPickler = compositePickler[SiteEvent]
    .addConcreteType[SiteCreated]
    .addConcreteType[SiteNameSet]
    .addConcreteType[SiteDescriptionSet]
    .addConcreteType[SiteRevisionSet]
    .addConcreteType[SitePublished]
    .addConcreteType[SiteUnpublished]
    .addConcreteType[SiteFlagAdded]
    .addConcreteType[SiteFlagRemoved]
    .addConcreteType[DomainAdded]
    .addConcreteType[DomainRemoved]
    .addConcreteType[PrimaryDomainSet]
    .addConcreteType[DefaultMetaTagAdded]
    .addConcreteType[DefaultMetaTagRemoved]
    .addConcreteType[PageAdded]
    .addConcreteType[PageRemoved]
    .addConcreteType[PageNameSet]
    .addConcreteType[PageMetaTagAdded]
    .addConcreteType[PageMetaTagRemoved]
    .addConcreteType[PageComponentAdded]
    .addConcreteType[PageComponentRemoved]
    .addConcreteType[PageComponentPositionSet]
    .addConcreteType[PageComponentPositionReset]
    .addConcreteType[TextComponentDataSet]
    .addConcreteType[ButtonComponentDataSet]
    .addConcreteType[BlogComponentDataSet]
    .addConcreteType[DomainEntryPointAdded]
    .addConcreteType[FreeEntryPointAdded]
    .addConcreteType[EntryPointRemoved]
    .addConcreteType[PrimaryEntryPointSet]
} 
Example 62
Source File: AppBase.scala    From suzaku   with Apache License 2.0 5 votes vote down vote up
package suzaku.app

import java.nio.ByteBuffer

import arteria.core._
import boopickle.DefaultBasic._
import suzaku.platform.{Logger, Transport}
import suzaku.ui.UIProtocol.UIChannel
import suzaku.ui.{Blueprint, UIManagerProxy, UIProtocol}
import suzaku.util.LoggerProtocol

abstract class AppBase(transport: Transport,
                       createHandler: (MessageChannelHandler[UIProtocol.type]) => AppRouterHandler = vm =>
                         new AppRouterHandler(vm))(implicit routerPickler: Pickler[RouterMessage] =
                                                     RouterMessage.defaultRouterPickler) {

  protected var logLevel = Logger.LogLevelDebug
  protected val logger: Logger = new Logger {
    import Logger._
    override def debug(message: => String): Unit = if (logLevel <= LogLevelDebug) println(message)
    override def info(message: => String): Unit  = if (logLevel <= LogLevelInfo) println(message)
    override def warn(message: => String): Unit  = if (logLevel <= LogLevelWarn) println(message)
    override def error(message: => String): Unit = if (logLevel <= LogLevelError) println(message)
  }
  protected val uiManager = new UIManagerProxy(logger, channelEstablished, flushMessages _)
  protected val router    = new MessageRouter[RouterMessage](createHandler(uiManager), false)

  // constructor
  // subscribe to messages from transport
  transport.subscribe(receive)
  // create channel for logger
  protected val loggerChannel = router.createChannel(LoggerProtocol)(MessageChannelHandler.empty[LoggerProtocol.type],
                                                                     LoggerProtocol.LoggerProtocolContext(),
                                                                     CreateLoggerChannel)

  // log about startup
  logger.info("Application starting")
  // send out first messages to establish router channel
  transport.send(router.flush())

  protected def main(): Unit

  protected def receive(data: ByteBuffer): Unit = {
    router.receive(data)
    // send pending messages
    if (router.hasPending) {
      transport.send(router.flush())
    }
  }

  protected def flushMessages(): Unit = {
    if (router.hasPending) {
      transport.send(router.flush())
    }
  }

  protected def channelEstablished(channel: UIChannel): Unit = {
    // logger.debug("UI channel established")
    // start application
    main()
  }

  def render(root: Blueprint): Unit = uiManager.render(root)
}

class AppRouterHandler(uiHandler: MessageChannelHandler[UIProtocol.type]) extends MessageRouterHandler[RouterMessage] {
  override def materializeChildChannel(id: Int,
    globalId: Int,
    router: MessageRouterBase,
    materializeChild: RouterMessage,
    contextReader: ChannelReader): Option[MessageChannelBase] = {
    materializeChild match {
      case CreateUIChannel =>
        val context = contextReader.read[UIProtocol.ChannelContext](UIProtocol.contextPickler)
        Some(new MessageChannel(UIProtocol)(id, globalId, router, uiHandler, context))
      case _ =>
        None
    }
  }
} 
Example 63
Source File: WebWorkerTransport.scala    From suzaku   with Apache License 2.0 5 votes vote down vote up
package suzaku.platform.web

import java.nio.ByteBuffer

import suzaku.platform.Transport
import org.scalajs.dom
import org.scalajs.dom.webworkers.{DedicatedWorkerGlobalScope, Worker}

import scala.scalajs.js
import scala.scalajs.js.typedarray.TypedArrayBufferOps._
import scala.scalajs.js.typedarray.{ArrayBuffer, TypedArrayBuffer}

abstract class WebWorkerTransport extends Transport {
  private val emptyHandler                    = (data: ByteBuffer) => ()
  private var dataHandler: ByteBuffer => Unit = emptyHandler

  override def subscribe(handler: ByteBuffer => Unit): () => Unit = {
    dataHandler = handler
    // return an unsubscribing function
    () =>
      {
        // restore our empty handler when called
        dataHandler = emptyHandler
      }
  }

  def receive(data: ArrayBuffer): Unit = dataHandler(TypedArrayBuffer.wrap(data))
}

class WorkerTransport(worker: DedicatedWorkerGlobalScope) extends WebWorkerTransport {
  override def send(data: ByteBuffer): Unit = {
    assert(data.hasTypedArray())
    val out = data.typedArray.buffer.slice(0, data.limit())
    worker.postMessage(out, js.Array(out).asInstanceOf[js.UndefOr[js.Array[dom.raw.Transferable]]])
  }
}

class WorkerClientTransport(worker: Worker) extends WebWorkerTransport {
  override def send(data: ByteBuffer): Unit = {
    assert(data.hasTypedArray())
    val out = data.typedArray.buffer.slice(0, data.limit())
    worker.postMessage(out, js.Array(out).asInstanceOf[js.UndefOr[js.Array[dom.raw.Transferable]]])
  }
} 
Example 64
Source File: TestUtils.scala    From iodb   with Creative Commons Zero v1.0 Universal 5 votes vote down vote up
package io.iohk.iodb

import java.io.File
import java.nio.ByteBuffer
import java.util.concurrent.{ExecutorService, TimeUnit}
import java.util.logging.Level

import scala.util.Random


  def runningTime[A](computation: => A): (Long, A) = {
    val s = System.currentTimeMillis()
    val res = computation
    (System.currentTimeMillis() - s, res)
  }

  def fromLong(id: Long): ByteArrayWrapper = {
    val b = ByteBuffer.allocate(8)
    b.putLong(0, id)
    ByteArrayWrapper(b.array())
  }


  def runnable(f: => Unit): Runnable =
    return () => {
      try {
        f
      } catch {
        case e: Throwable => {
          Utils.LOG.log(Level.SEVERE, "Background task failed", e)
        }
      }
    }

  def waitForFinish(exec: ExecutorService): Unit = {
    exec.shutdown()
    exec.awaitTermination(400, TimeUnit.DAYS)
  }


  def withTempDir(ff: (File) => Unit) {
    val iFile = TestUtils.tempDir()
    try {
      ff(iFile)
    } finally {
      TestUtils.deleteRecur(iFile)
    }
  }
} 
Example 65
Source File: NsUUID.scala    From seals   with Apache License 2.0 5 votes vote down vote up
package dev.tauri.seals
package core

import java.util.UUID
import java.security.MessageDigest
import java.nio.{ ByteBuffer, Buffer }
import java.nio.charset.StandardCharsets

import scodec.bits.ByteVector

private[seals] object NsUUID {

  def uuid5(ns: UUID, name: String): UUID =
    uuid5bytes(ns, ByteBuffer.wrap(name.getBytes(StandardCharsets.UTF_8)))

  def uuid5bv(ns: UUID, name: ByteVector): UUID =
    uuid5bytes(ns, name.toByteBuffer)

  private def uuid5bytes(ns: UUID, name: ByteBuffer): UUID = {
    val buf = ByteBuffer.allocate(16) // network byte order by default
    putUUIDToBuf(ns, buf)
    (buf : Buffer).rewind()
    val h = sha1()
    h.update(buf)
    h.update(name)
    val arr: Array[Byte] = h.digest().take(16)
    arr(6) = (arr(6) & 0x0f).toByte // clear version
    arr(6) = (arr(6) | 0x50).toByte // version 5
    arr(8) = (arr(8) & 0x3f).toByte // clear variant
    arr(8) = (arr(8) | 0x80).toByte // variant RFC4122
    (buf : Buffer).rewind()
    buf.put(arr)
    (buf : Buffer).rewind()
    val msl = buf.getLong()
    val lsl = buf.getLong()
    new UUID(msl, lsl)
  }

  def uuid5nested(root: UUID, names: String*): UUID =
    names.foldLeft(root)(uuid5)

  def uuid5nestedBv(root: UUID, names: ByteVector*): UUID = {
    val buf = ByteVector.concat(names).toByteBuffer
    uuid5bytes(root, buf)
  }

  def uuid5nestedNsNm(name: String, ns1: UUID, nss: UUID*): UUID =
    uuid5(uuid5nestedNs(ns1, nss: _*), name)

  def uuid5nestedNs(ns1: UUID, nss: UUID*): UUID = {
    val buf = ByteBuffer.allocate(16)
    nss.foldLeft(ns1) { (st, u) =>
      putUUIDToBuf(u, buf)
      (buf : Buffer).rewind()
      val r = uuid5bytes(st, buf)
      (buf : Buffer).rewind()
      r
    }
  }

  private def putUUIDToBuf(u: UUID, buf: ByteBuffer): Unit = {
    buf.putLong(u.getMostSignificantBits)
    buf.putLong(u.getLeastSignificantBits)
  }

  def bvFromUUID(u: UUID): ByteVector = {
    val buf = ByteBuffer.allocate(16)
    putUUIDToBuf(u, buf)
    (buf : Buffer).rewind()
    ByteVector.view(buf)
  }

  private def sha1(): MessageDigest =
    MessageDigest.getInstance("SHA-1")
} 
Example 66
Source File: Json4sSerializer.scala    From reliable-http-client   with Apache License 2.0 5 votes vote down vote up
package rhttpc.akkapersistence.json4s

import java.nio.ByteBuffer
import java.nio.charset.Charset

import akka.actor.ExtendedActorSystem
import akka.serialization.Serializer
import org.json4s.native.Serialization._
import org.json4s.{DefaultFormats, Formats, TypeHints}
import rhttpc.transport.json4s.{AllTypeHints, ObjectSerializer}

class Json4sSerializer(system: ExtendedActorSystem) extends Serializer {
  import Json4sSerializer._
  import rhttpc.transport.json4s.CommonFormats._

  override def identifier: Int = ID

  override def includeManifest: Boolean = true

  override def fromBinary(bytes: Array[Byte], manifestOpt: Option[Class[_]]): AnyRef = {
    implicit val manifest = manifestOpt match {
      case Some(x) => Manifest.classType(x)
      case None    => Manifest.AnyRef
    }
    read(new String(bytes, UTF8))
  }

  override def toBinary(o: AnyRef): Array[Byte] = {
    writePretty(o).getBytes(UTF8)
  }
}

object Json4sSerializer {
  private val UTF8: Charset = Charset.forName("UTF-8")
  private val ID: Int = ByteBuffer.wrap("json4s".getBytes(UTF8)).getInt
} 
Example 67
Source File: KuduValueWriter.scala    From eel-sdk   with Apache License 2.0 5 votes vote down vote up
package io.eels.component.kudu

import java.nio.ByteBuffer

import io.eels.coercion.{DoubleCoercer, FloatCoercer, IntCoercer, LongCoercer}
import org.apache.kudu.client.PartialRow

trait KuduValueWriter {
  def write(row: PartialRow, index: Int, value: Any): Unit
}

object KuduStringWriter extends KuduValueWriter {
  override def write(row: PartialRow, index: Int, value: Any): Unit = {
    row.addString(index, value.toString)
  }
}

object KuduBinaryWriter extends KuduValueWriter {
  override def write(row: PartialRow, index: Int, value: Any): Unit = {
    value match {
      case array: Array[Byte] => row.addBinary(index, array)
      case buffer: ByteBuffer => row.addBinary(index, buffer)
    }
  }
}

object KuduLongWriter extends KuduValueWriter {
  override def write(row: PartialRow, index: Int, value: Any): Unit = {
    row.addLong(index, LongCoercer.coerce(value))
  }
}

object KuduIntWriter extends KuduValueWriter {
  override def write(row: PartialRow, index: Int, value: Any): Unit = {
    row.addInt(index, IntCoercer.coerce(value))
  }
}

object KuduBooleanWriter extends KuduValueWriter {
  override def write(row: PartialRow, index: Int, value: Any): Unit = {
    row.addBoolean(index, value.asInstanceOf[Boolean])
  }
}

object KuduByteWriter extends KuduValueWriter {
  override def write(row: PartialRow, index: Int, value: Any): Unit = {
    row.addByte(index, value.asInstanceOf[Byte])
  }
}

object KuduDoubleWriter extends KuduValueWriter {
  override def write(row: PartialRow, index: Int, value: Any): Unit = {
    row.addDouble(index, DoubleCoercer.coerce(value))
  }
}

object KuduFloatWriter extends KuduValueWriter {
  override def write(row: PartialRow, index: Int, value: Any): Unit = {
    row.addFloat(index, FloatCoercer.coerce(value))
  }
}

object KuduShortWriter extends KuduValueWriter {
  override def write(row: PartialRow, index: Int, value: Any): Unit = {
    row.addShort(index, value.asInstanceOf[Short])
  }
} 
Example 68
Source File: LeveldbNumericIdentifierStore.scala    From eventuate   with Apache License 2.0 5 votes vote down vote up
package com.rbmhtechnology.eventuate.log.leveldb

import java.nio.ByteBuffer

import org.iq80.leveldb.{ DB, DBIterator }

private class LeveldbNumericIdentifierStore(leveldb: DB, classifier: Int) {
  private var idMap: Map[String, Int] =
    Map.empty

  private val idKeyEnd: Int =
    Int.MaxValue

  private val idKeyEndBytes: Array[Byte] =
    idKeyBytes(idKeyEnd)

  leveldb.put(idKeyEndBytes, Array.empty[Byte])

  def numericId(id: String): Int = idMap.get(id) match {
    case None    => writeIdMapping(id, idMap.size + 1)
    case Some(v) => v
  }

  def findId(numericId: Int): Option[String] =
    idMap.find { case (_, nid) => nid == numericId }.map(_._1)

  def readIdMap(iter: DBIterator): Unit = {
    iter.seek(idKeyBytes(0))
    idMap = readIdMap(Map.empty, iter)
  }

  private def readIdMap(idMap: Map[String, Int], iter: DBIterator): Map[String, Int] = {
    if (!iter.hasNext) idMap else {
      val nextEntry = iter.next()
      val nextKey = idKey(nextEntry.getKey)
      if (nextKey == idKeyEnd) idMap else {
        val nextVal = new String(nextEntry.getValue, "UTF-8")
        readIdMap(idMap + (nextVal -> nextKey), iter)
      }
    }
  }

  private def writeIdMapping(id: String, nid: Int): Int = {
    idMap = idMap + (id -> nid)
    leveldb.put(idKeyBytes(nid), id.getBytes("UTF-8"))
    nid
  }

  private def idKeyBytes(nid: Int): Array[Byte] = {
    val bb = ByteBuffer.allocate(8)
    bb.putInt(classifier)
    bb.putInt(nid)
    bb.array
  }

  private def idKey(a: Array[Byte]): Int = {
    val bb = ByteBuffer.wrap(a)
    bb.getInt
    bb.getInt
  }
} 
Example 69
Source File: LeveldbDeletionMetadataStore.scala    From eventuate   with Apache License 2.0 5 votes vote down vote up
package com.rbmhtechnology.eventuate.log.leveldb

import java.nio.ByteBuffer

import com.rbmhtechnology.eventuate.log.DeletionMetadata
import com.rbmhtechnology.eventuate.log.leveldb.LeveldbEventLog.WithBatch
import com.rbmhtechnology.eventuate.log.leveldb.LeveldbEventLog.longBytes
import org.iq80.leveldb.DB
import org.iq80.leveldb.WriteOptions

private class LeveldbDeletionMetadataStore(val leveldb: DB, val leveldbWriteOptions: WriteOptions, classifier: Int) extends WithBatch {
  private val DeletedToSequenceNrKey: Int = 1
  private val RemoteLogIdsKey: Int = 2

  private val StringSetSeparatorChar = '\u0000'

  def writeDeletionMetadata(info: DeletionMetadata): Unit = withBatch { batch =>
    batch.put(idKeyBytes(DeletedToSequenceNrKey), longBytes(info.toSequenceNr))
    batch.put(idKeyBytes(RemoteLogIdsKey), stringSetBytes(info.remoteLogIds))
  }

  def readDeletionMetadata(): DeletionMetadata = {
    val toSequenceNr = longFromBytes(leveldb.get(idKeyBytes(DeletedToSequenceNrKey)))
    val remoteLogIds = stringSetFromBytes(leveldb.get(idKeyBytes(RemoteLogIdsKey)))
    DeletionMetadata(toSequenceNr, remoteLogIds)
  }

  private def idKeyBytes(key: Int): Array[Byte] = {
    val bb = ByteBuffer.allocate(8)
    bb.putInt(classifier)
    bb.putInt(key)
    bb.array
  }

  private def longFromBytes(longBytes: Array[Byte]): Long =
    if (longBytes == null) 0 else LeveldbEventLog.longFromBytes(longBytes)

  private def stringSetBytes(set: Set[String]): Array[Byte] =
    set.mkString(StringSetSeparatorChar.toString).getBytes("UTF-8")

  private def stringSetFromBytes(setBytes: Array[Byte]): Set[String] =
    if (setBytes == null || setBytes.length == 0)
      Set.empty
    else
      new String(setBytes, "UTF-8").split(StringSetSeparatorChar).toSet
} 
Example 70
Source File: LeveldbReplicationProgressStore.scala    From eventuate   with Apache License 2.0 5 votes vote down vote up
package com.rbmhtechnology.eventuate.log.leveldb

import java.nio.ByteBuffer

import org.iq80.leveldb.{ DB, DBIterator, WriteBatch }

private class LeveldbReplicationProgressStore(leveldb: DB, classifier: Int, numericId: String => Int, findId: Int => Option[String]) {
  private val rpKeyEnd: Int =
    Int.MaxValue

  private val rpKeyEndBytes: Array[Byte] =
    rpKeyBytes(rpKeyEnd)

  leveldb.put(rpKeyEndBytes, Array.empty[Byte])

  def writeReplicationProgress(logId: String, logSnr: Long, batch: WriteBatch): Unit = {
    val nid = numericId(logId)
    batch.put(rpKeyBytes(nid), LeveldbEventLog.longBytes(logSnr))
  }

  def readReplicationProgress(logId: String): Long = {
    val nid = numericId(logId)
    val progress = leveldb.get(rpKeyBytes(nid))
    if (progress == null) 0L else LeveldbEventLog.longFromBytes(progress)
  }

  def readReplicationProgresses(iter: DBIterator): Map[String, Long] = {
    iter.seek(rpKeyBytes(0))
    readReplicationProgresses(Map.empty, iter).foldLeft(Map.empty[String, Long]) {
      case (acc, (nid, progress)) => findId(nid) match {
        case Some(id) => acc + (id -> progress)
        case None     => acc
      }
    }
  }

  private def readReplicationProgresses(rpMap: Map[Int, Long], iter: DBIterator): Map[Int, Long] = {
    if (!iter.hasNext) rpMap else {
      val nextEntry = iter.next()
      val nextKey = rpKey(nextEntry.getKey)
      if (nextKey == rpKeyEnd) rpMap else {
        val nextVal = LeveldbEventLog.longFromBytes(nextEntry.getValue)
        readReplicationProgresses(rpMap + (nextKey -> nextVal), iter)
      }
    }
  }

  private def rpKeyBytes(nid: Int): Array[Byte] = {
    val bb = ByteBuffer.allocate(8)
    bb.putInt(classifier)
    bb.putInt(nid)
    bb.array
  }

  private def rpKey(a: Array[Byte]): Int = {
    val bb = ByteBuffer.wrap(a)
    bb.getInt
    bb.getInt
  }
} 
Example 71
Source File: Chapter10.scala    From Learning-Spark-SQL   with MIT License 5 votes vote down vote up
//Code for Chapter 10 to be executed in Spark shell. For all other code from the BigDL library follow the instructions and commands in the book.
//Note that the code in this Chapter uses Spark 2.1 due to some bugs.

//Execute the following on the command prompt to start the Spark shell
source /Users/aurobindosarkar/Downloads/BigDL-master/scripts/bigdl.sh
Aurobindos-MacBook-Pro-2:spark-2.1.0-bin-hadoop2.7 aurobindosarkar$ bin/spark-shell --properties-file /Users/aurobindosarkar/Downloads/BigDL-master/spark/dist/target/bigdl-0.2.0-SNAPSHOT-spark-2.0.0-scala-2.11.8-mac-dist/conf/spark-bigdl.conf --jars /Users/aurobindosarkar/Downloads/BigDL-master/spark/dist/target/bigdl-0.2.0-SNAPSHOT-spark-2.0.0-scala-2.11.8-mac-dist/lib/bigdl-0.2.0-SNAPSHOT-jar-with-dependencies.jar

import com.intel.analytics.bigdl._
import com.intel.analytics.bigdl.dataset.DataSet
import com.intel.analytics.bigdl.dataset.image.{BytesToGreyImg, GreyImgNormalizer, GreyImgToBatch, GreyImgToSample}
import com.intel.analytics.bigdl.nn.{ClassNLLCriterion, Module}
import com.intel.analytics.bigdl.numeric.NumericFloat
import com.intel.analytics.bigdl.optim._
import com.intel.analytics.bigdl.utils.{Engine, LoggerFilter, T, Table}
import com.intel.analytics.bigdl.nn._
import java.nio.ByteBuffer
import java.nio.file.{Files, Path, Paths}
import com.intel.analytics.bigdl.dataset.ByteRecord
import com.intel.analytics.bigdl.utils.File

val trainData = "/Users/aurobindosarkar/Downloads/mnist/train-images-idx3-ubyte"
val trainLabel = "/Users/aurobindosarkar/Downloads/mnist/train-labels-idx1-ubyte"
val validationData = "/Users/aurobindosarkar/Downloads/mnist/t10k-images-idx3-ubyte"
val validationLabel = "/Users/aurobindosarkar/Downloads/mnist/t10k-labels-idx1-ubyte"
val nodeNumber = 1
val coreNumber = 2
Engine.init
val model = Sequential[Float]()
val classNum = 10
val batchSize = 12
model.add(Reshape(Array(1, 28, 28))).add(SpatialConvolution(1, 6, 5, 5)).add(Tanh()).add(SpatialMaxPooling(2, 2, 2, 2)).add(Tanh()).add(SpatialConvolution(6, 12, 5, 5)).add(SpatialMaxPooling(2, 2, 2, 2)).add(Reshape(Array(12 * 4 * 4))).add(Linear(12 * 4 * 4, 100)).add(Tanh()).add(Linear(100, classNum)).add(LogSoftMax())

def load(featureFile: String, labelFile: String): Array[ByteRecord] = {
    val featureBuffer = ByteBuffer.wrap(Files.readAllBytes(Paths.get(featureFile)))
    val labelBuffer = ByteBuffer.wrap(Files.readAllBytes(Paths.get(labelFile)));
    val labelMagicNumber = labelBuffer.getInt();
    require(labelMagicNumber == 2049);
    val featureMagicNumber = featureBuffer.getInt();
    require(featureMagicNumber == 2051);
    val labelCount = labelBuffer.getInt();
    val featureCount = featureBuffer.getInt();
    require(labelCount == featureCount);
    val rowNum = featureBuffer.getInt();
    val colNum = featureBuffer.getInt();
    val result = new Array[ByteRecord](featureCount);
    var i = 0;
    while (i < featureCount) {
      val img = new Array[Byte]((rowNum * colNum));
      var y = 0;
      while (y < rowNum) {
        var x = 0;
        while (x < colNum) {
          img(x + y * colNum) = featureBuffer.get();
          x += 1;
        }
        y += 1;
      }
      result(i) = ByteRecord(img, labelBuffer.get().toFloat + 1.0f);
      i += 1;
    }
    result;
  }
val trainMean = 0.13066047740239506
val trainStd = 0.3081078
val trainSet = DataSet.array(load(trainData, trainLabel), sc) -> BytesToGreyImg(28, 28) -> GreyImgNormalizer(trainMean, trainStd) -> GreyImgToBatch(batchSize)
val optimizer = Optimizer(model = model, dataset = trainSet, criterion = ClassNLLCriterion[Float]())   
val testMean = 0.13251460696903547
val testStd = 0.31048024
val maxEpoch = 2
val validationSet = DataSet.array(load(validationData, validationLabel), sc) -> BytesToGreyImg(28, 28) -> GreyImgNormalizer(testMean, testStd) -> GreyImgToBatch(batchSize)
optimizer.setEndWhen(Trigger.maxEpoch(2))
optimizer.setState(T("learningRate" -> 0.05, "learningRateDecay" -> 0.0))
optimizer.setCheckpoint("/Users/aurobindosarkar/Downloads/mnist/checkpoint", Trigger.severalIteration(500))
optimizer.setValidation(trigger = Trigger.everyEpoch, dataset = validationSet, vMethods = Array(new Top1Accuracy, new Top5Accuracy[Float], new Loss[Float]))
optimizer.optimize()
model.save("/Users/aurobindosarkar/Downloads/mnist/model")
val model = Module.load[Float]("/Users/aurobindosarkar/Downloads/mnist/model")
val partitionNum = 2
val rddData = sc.parallelize(load(validationData, validationLabel), partitionNum)
val transformer = BytesToGreyImg(28, 28) -> GreyImgNormalizer(testMean, testStd) -> GreyImgToSample()
val evaluationSet = transformer(rddData)
val result = model.evaluate(evaluationSet, Array(new Top1Accuracy[Float]), Some(batchSize))
result.foreach(r => println(s"${r._2} is ${r._1}")) 
Example 72
Source File: JdkHttpClient.scala    From http4s-jdk-http-client   with Apache License 2.0 5 votes vote down vote up
package org.http4s.client.jdkhttpclient

import java.net.URI
import java.net.http.HttpRequest.BodyPublishers
import java.net.http.HttpResponse.BodyHandlers
import java.net.http.{HttpClient, HttpRequest, HttpResponse}
import java.nio.ByteBuffer
import java.util
import java.util.concurrent.Flow

import cats.ApplicativeError
import cats.effect._
import cats.implicits._
import fs2.concurrent.SignallingRef
import fs2.interop.reactivestreams._
import fs2.{Chunk, Stream}
import org.http4s.client.Client
import org.http4s.client.jdkhttpclient.compat.CollectionConverters._
import org.http4s.internal.fromCompletionStage
import org.http4s.util.CaseInsensitiveString
import org.http4s.{Header, Headers, HttpVersion, Request, Response, Status}
import org.reactivestreams.FlowAdapters

object JdkHttpClient {

  
  def simple[F[_]](implicit F: ConcurrentEffect[F], CS: ContextShift[F]): F[Client[F]] =
    F.delay(HttpClient.newHttpClient()).map(apply(_))

  def convertHttpVersionFromHttp4s[F[_]](
      version: HttpVersion
  )(implicit F: ApplicativeError[F, Throwable]): F[HttpClient.Version] =
    version match {
      case HttpVersion.`HTTP/1.1` => HttpClient.Version.HTTP_1_1.pure[F]
      case HttpVersion.`HTTP/2.0` => HttpClient.Version.HTTP_2.pure[F]
      case _ => F.raiseError(new IllegalArgumentException("invalid HTTP version"))
    }

  // see jdk.internal.net.http.common.Utils#DISALLOWED_HEADERS_SET
  private val restrictedHeaders =
    Set(
      "connection",
      "content-length",
      "date",
      "expect",
      "from",
      "host",
      "upgrade",
      "via",
      "warning"
    ).map(CaseInsensitiveString(_))
} 
Example 73
Source File: ManagerUtilsSuite.scala    From darwin   with Apache License 2.0 5 votes vote down vote up
package it.agilelab.darwin.app.mock

import java.nio.{ByteBuffer, ByteOrder}

import com.typesafe.config.ConfigFactory
import it.agilelab.darwin.manager.AvroSchemaManagerFactory
import it.agilelab.darwin.manager.util.{AvroSingleObjectEncodingUtils, ConfigurationKeys}
import it.agilelab.darwin.manager.util.ByteArrayUtils._

import scala.util.Random
import org.scalatest.flatspec.AnyFlatSpec
import org.scalatest.matchers.should.Matchers

class BigEndianManagerUtilsSuite extends ManagerUtilsSuite(ByteOrder.BIG_ENDIAN)

class LittleEndianManagerUtilsSuite extends ManagerUtilsSuite(ByteOrder.LITTLE_ENDIAN)

abstract class ManagerUtilsSuite(endianness: ByteOrder) extends AnyFlatSpec with Matchers {

  "AvroSchemaManager utilities" should "create a Single-Object encoded byte array" in {
    val ORIGINAL_LENGTH: Int = 10
    val originalSchema = SchemaReader.readFromResources("OneField.avsc")
    val config =
      ConfigFactory
        .parseMap(new java.util.HashMap[String, String]() {
          {
            put(ConfigurationKeys.MANAGER_TYPE, ConfigurationKeys.CACHED_EAGER)
            put(ConfigurationKeys.ENDIANNESS, endianness.toString)
          }
        })
        .withFallback(ConfigFactory.load())
        .resolve()
    val manager = AvroSchemaManagerFactory.initialize(config)
    manager.registerAll(Seq(originalSchema))
    val originalPayload = new Array[Byte](ORIGINAL_LENGTH)
    Random.nextBytes(originalPayload)
    val data: Array[Byte] = manager.generateAvroSingleObjectEncoded(originalPayload, originalSchema)
    assert(AvroSingleObjectEncodingUtils.isAvroSingleObjectEncoded(data))
    val (schema, payload) = manager.retrieveSchemaAndAvroPayload(data)
    assert(schema == originalSchema)
    assert(originalPayload sameElements payload)
  }

  it should "convert a long to byte array and back" in {
    val longs = (1 to 10).map(_ => Random.nextLong())

    assert(
      longs == longs.map(
        x =>
          AvroSingleObjectEncodingUtils
            .readLong(ByteBuffer.wrap(x.longToByteArray(endianness)), endianness)
      )
    )
  }

} 
Example 74
Source File: ByteArrayUtils.scala    From darwin   with Apache License 2.0 5 votes vote down vote up
package it.agilelab.darwin.manager.util

import java.io.OutputStream
import java.nio.{ByteBuffer, ByteOrder}

import it.agilelab.darwin.common.LONG_SIZE


private[darwin] object ByteArrayUtils {

  implicit class EnrichedLong(val l: Long) extends AnyVal {
    
    def writeToStream(os: OutputStream, endianness: ByteOrder): Unit = {
      endianness match {
        case ByteOrder.BIG_ENDIAN =>
          os.write((l >>> 56).asInstanceOf[Int])
          os.write((l >>> 48).asInstanceOf[Int])
          os.write((l >>> 40).asInstanceOf[Int])
          os.write((l >>> 32).asInstanceOf[Int])
          os.write((l >>> 24).asInstanceOf[Int])
          os.write((l >>> 16).asInstanceOf[Int])
          os.write((l >>> 8).asInstanceOf[Int])
          os.write((l >>> 0).asInstanceOf[Int])
        case ByteOrder.LITTLE_ENDIAN =>
          os.write((l >>> 0).asInstanceOf[Int])
          os.write((l >>> 8).asInstanceOf[Int])
          os.write((l >>> 16).asInstanceOf[Int])
          os.write((l >>> 24).asInstanceOf[Int])
          os.write((l >>> 32).asInstanceOf[Int])
          os.write((l >>> 40).asInstanceOf[Int])
          os.write((l >>> 48).asInstanceOf[Int])
          os.write((l >>> 56).asInstanceOf[Int])
      }
    }
  }

  def arrayEquals(b1: Array[Byte], b2: Array[Byte], start1: Int, start2: Int, length: Int): Boolean = {
    require(length > 0, "length must be positive")
    var i = start1
    var j = start2
    var areEqual = true
    while (areEqual && i < start1 + length) {
      if (b1(i) != b2(j)) {
        areEqual = false
      }
      i += 1
      j += 1
    }
    areEqual
  }


} 
Example 75
Source File: DefaultBodyWritables.scala    From play-ws   with Apache License 2.0 5 votes vote down vote up
package play.api.libs.ws

import java.io.File
import java.nio.ByteBuffer
import java.util.function.Supplier

import akka.stream.scaladsl.StreamConverters.fromInputStream
import akka.stream.scaladsl.FileIO
import akka.stream.scaladsl.Source
import akka.util.ByteString

import scala.compat.java8.FunctionConverters.asScalaFromSupplier


  implicit val writeableOf_urlEncodedForm: BodyWritable[Map[String, Seq[String]]] = {
    import java.net.URLEncoder
    BodyWritable(
      formData =>
        InMemoryBody(
          ByteString.fromString(
            formData.flatMap(item => item._2.map(c => s"${item._1}=${URLEncoder.encode(c, "UTF-8")}")).mkString("&")
          )
        ),
      "application/x-www-form-urlencoded"
    )
  }

  implicit val writeableOf_urlEncodedSimpleForm: BodyWritable[Map[String, String]] = {
    writeableOf_urlEncodedForm.map[Map[String, String]](_.map(kv => kv._1 -> Seq(kv._2)))
  }

}

object DefaultBodyWritables extends DefaultBodyWritables 
Example 76
Source File: NullableColumnAccessor.scala    From XSQL   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.sql.execution.columnar

import java.nio.{ByteBuffer, ByteOrder}

import org.apache.spark.sql.catalyst.InternalRow

private[columnar] trait NullableColumnAccessor extends ColumnAccessor {
  private var nullsBuffer: ByteBuffer = _
  private var nullCount: Int = _
  private var seenNulls: Int = 0

  private var nextNullIndex: Int = _
  private var pos: Int = 0

  abstract override protected def initialize(): Unit = {
    nullsBuffer = underlyingBuffer.duplicate().order(ByteOrder.nativeOrder())
    nullCount = ByteBufferHelper.getInt(nullsBuffer)
    nextNullIndex = if (nullCount > 0) ByteBufferHelper.getInt(nullsBuffer) else -1
    pos = 0

    underlyingBuffer.position(underlyingBuffer.position() + 4 + nullCount * 4)
    super.initialize()
  }

  abstract override def extractTo(row: InternalRow, ordinal: Int): Unit = {
    if (pos == nextNullIndex) {
      seenNulls += 1

      if (seenNulls < nullCount) {
        nextNullIndex = ByteBufferHelper.getInt(nullsBuffer)
      }

      row.setNullAt(ordinal)
    } else {
      super.extractTo(row, ordinal)
    }

    pos += 1
  }

  abstract override def hasNext: Boolean = seenNulls < nullCount || super.hasNext
} 
Example 77
Source File: NullableColumnBuilder.scala    From XSQL   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.sql.execution.columnar

import java.nio.{ByteBuffer, ByteOrder}

import org.apache.spark.sql.catalyst.InternalRow


private[columnar] trait NullableColumnBuilder extends ColumnBuilder {
  protected var nulls: ByteBuffer = _
  protected var nullCount: Int = _
  private var pos: Int = _

  abstract override def initialize(
      initialSize: Int,
      columnName: String,
      useCompression: Boolean): Unit = {

    nulls = ByteBuffer.allocate(1024)
    nulls.order(ByteOrder.nativeOrder())
    pos = 0
    nullCount = 0
    super.initialize(initialSize, columnName, useCompression)
  }

  abstract override def appendFrom(row: InternalRow, ordinal: Int): Unit = {
    columnStats.gatherStats(row, ordinal)
    if (row.isNullAt(ordinal)) {
      nulls = ColumnBuilder.ensureFreeSpace(nulls, 4)
      nulls.putInt(pos)
      nullCount += 1
    } else {
      super.appendFrom(row, ordinal)
    }
    pos += 1
  }

  abstract override def build(): ByteBuffer = {
    val nonNulls = super.build()
    val nullDataLen = nulls.position()

    nulls.limit(nullDataLen)
    nulls.rewind()

    val buffer = ByteBuffer
      .allocate(4 + nullDataLen + nonNulls.remaining())
      .order(ByteOrder.nativeOrder())
      .putInt(nullCount)
      .put(nulls)
      .put(nonNulls)

    buffer.rewind()
    buffer
  }

  protected def buildNonNulls(): ByteBuffer = {
    nulls.limit(nulls.position()).rewind()
    super.build()
  }
} 
Example 78
Source File: CompressibleColumnBuilder.scala    From XSQL   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.sql.execution.columnar.compression

import java.nio.{ByteBuffer, ByteOrder}

import org.apache.spark.internal.Logging
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.execution.columnar.{ColumnBuilder, NativeColumnBuilder}
import org.apache.spark.sql.types.AtomicType
import org.apache.spark.unsafe.Platform


private[columnar] trait CompressibleColumnBuilder[T <: AtomicType]
  extends ColumnBuilder with Logging {

  this: NativeColumnBuilder[T] with WithCompressionSchemes =>

  var compressionEncoders: Seq[Encoder[T]] = _

  abstract override def initialize(
      initialSize: Int,
      columnName: String,
      useCompression: Boolean): Unit = {

    compressionEncoders =
      if (useCompression) {
        schemes.filter(_.supports(columnType)).map(_.encoder[T](columnType))
      } else {
        Seq(PassThrough.encoder(columnType))
      }
    super.initialize(initialSize, columnName, useCompression)
  }

  // The various compression schemes, while saving memory use, cause all of the data within
  // the row to become unaligned, thus causing crashes.  Until a way of fixing the compression
  // is found to also allow aligned accesses this must be disabled for SPARC.

  protected def isWorthCompressing(encoder: Encoder[T]) = {
    CompressibleColumnBuilder.unaligned && encoder.compressionRatio < 0.8
  }

  private def gatherCompressibilityStats(row: InternalRow, ordinal: Int): Unit = {
    compressionEncoders.foreach(_.gatherCompressibilityStats(row, ordinal))
  }

  abstract override def appendFrom(row: InternalRow, ordinal: Int): Unit = {
    super.appendFrom(row, ordinal)
    if (!row.isNullAt(ordinal)) {
      gatherCompressibilityStats(row, ordinal)
    }
  }

  override def build(): ByteBuffer = {
    val nonNullBuffer = buildNonNulls()
    val encoder: Encoder[T] = {
      val candidate = compressionEncoders.minBy(_.compressionRatio)
      if (isWorthCompressing(candidate)) candidate else PassThrough.encoder(columnType)
    }

    // Header = null count + null positions
    val headerSize = 4 + nulls.limit()
    val compressedSize = if (encoder.compressedSize == 0) {
      nonNullBuffer.remaining()
    } else {
      encoder.compressedSize
    }

    val compressedBuffer = ByteBuffer
      // Reserves 4 bytes for compression scheme ID
      .allocate(headerSize + 4 + compressedSize)
      .order(ByteOrder.nativeOrder)
      // Write the header
      .putInt(nullCount)
      .put(nulls)

    logDebug(s"Compressor for [$columnName]: $encoder, ratio: ${encoder.compressionRatio}")
    encoder.compress(nonNullBuffer, compressedBuffer)
  }
}

private[columnar] object CompressibleColumnBuilder {
  val unaligned = Platform.unaligned()
} 
Example 79
Source File: NullableColumnAccessorSuite.scala    From XSQL   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.sql.execution.columnar

import java.nio.ByteBuffer

import org.apache.spark.SparkFunSuite
import org.apache.spark.sql.catalyst.CatalystTypeConverters
import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeProjection}
import org.apache.spark.sql.types._

class TestNullableColumnAccessor[JvmType](
    buffer: ByteBuffer,
    columnType: ColumnType[JvmType])
  extends BasicColumnAccessor(buffer, columnType)
  with NullableColumnAccessor

object TestNullableColumnAccessor {
  def apply[JvmType](buffer: ByteBuffer, columnType: ColumnType[JvmType])
    : TestNullableColumnAccessor[JvmType] = {
    new TestNullableColumnAccessor(buffer, columnType)
  }
}

class NullableColumnAccessorSuite extends SparkFunSuite {
  import org.apache.spark.sql.execution.columnar.ColumnarTestUtils._

  Seq(
    NULL, BOOLEAN, BYTE, SHORT, INT, LONG, FLOAT, DOUBLE,
    STRING, BINARY, COMPACT_DECIMAL(15, 10), LARGE_DECIMAL(20, 10),
    STRUCT(StructType(StructField("a", StringType) :: Nil)),
    ARRAY(ArrayType(IntegerType)), MAP(MapType(IntegerType, StringType)))
    .foreach {
    testNullableColumnAccessor(_)
  }

  def testNullableColumnAccessor[JvmType](
      columnType: ColumnType[JvmType]): Unit = {

    val typeName = columnType.getClass.getSimpleName.stripSuffix("$")
    val nullRow = makeNullRow(1)

    test(s"Nullable $typeName column accessor: empty column") {
      val builder = TestNullableColumnBuilder(columnType)
      val accessor = TestNullableColumnAccessor(builder.build(), columnType)
      assert(!accessor.hasNext)
    }

    test(s"Nullable $typeName column accessor: access null values") {
      val builder = TestNullableColumnBuilder(columnType)
      val randomRow = makeRandomRow(columnType)
      val proj = UnsafeProjection.create(Array[DataType](columnType.dataType))

      (0 until 4).foreach { _ =>
        builder.appendFrom(proj(randomRow), 0)
        builder.appendFrom(proj(nullRow), 0)
      }

      val accessor = TestNullableColumnAccessor(builder.build(), columnType)
      val row = new GenericInternalRow(1)
      val converter = CatalystTypeConverters.createToScalaConverter(columnType.dataType)

      (0 until 4).foreach { _ =>
        assert(accessor.hasNext)
        accessor.extractTo(row, 0)
        assert(converter(row.get(0, columnType.dataType))
          === converter(randomRow.get(0, columnType.dataType)))

        assert(accessor.hasNext)
        accessor.extractTo(row, 0)
        assert(row.isNullAt(0))
      }

      assert(!accessor.hasNext)
    }
  }
} 
Example 80
Source File: TruckDataScheme.scala    From trucking-iot   with Apache License 2.0 5 votes vote down vote up
package com.orendainx.trucking.storm.schemes

import java.nio.ByteBuffer

import org.apache.storm.tuple.{Fields, Values}


object TruckDataScheme extends DelimitedScheme("\\|") {

  override def deserialize(buffer: ByteBuffer): Values = {

    // Extract data from buffer
    val strings = deserializeStringAndSplit(buffer)
    val eventTime = strings(0)
    val truckId = strings(1)
    val driverId = strings(2)
    val driverName = strings(3)
    val routeId = strings(4)
    val routeName = strings(5)
    val latitude = strings(6)
    val longitude = strings(7)
    val speed = strings(8)
    val eventType = strings(9)

    // TODO: Q: Feed strings directly into Values()?  Benefit to unpackaging byteBuffer into each field's appropriate type?
    new Values(eventTime, truckId, driverId, driverName, routeId, routeName, latitude, longitude, speed, eventType)
  }

  override def getOutputFields: Fields =
    new Fields("eventTime", "truckId", "driverId", "driverName", "routeId", "routeName", "latitude", "longitude", "speed", "eventType")
} 
Example 81
Source File: MerkleDigest.scala    From JustinDB   with Apache License 2.0 5 votes vote down vote up
package justin.db.merkletrees

trait MerkleDigest[T] {
  def digest(t: T): Digest
}

object MerkleDigest {

  implicit object CRC32 extends MerkleDigest[Block] {
    import java.nio.ByteBuffer
    import java.util.zip.CRC32

    override def digest(t: Block): Digest = {
      val digest = new CRC32()
      digest.update(t)

      val buffer = ByteBuffer.allocate(8)
      buffer.putLong(digest.getValue)

      Digest(buffer.array())
    }
  }

  implicit object MD5 extends MerkleDigest[Block] {
    import java.security.MessageDigest

    override def digest(t: Block): Digest = {
      Digest(MessageDigest.getInstance("MD5").digest(t))
    }
  }
}

case class Digest(hash: Array[Byte]) extends AnyVal {
  def +(that: Digest): Digest = Digest(this.hash ++ that.hash)
  def ==(that: Digest): Boolean = this.hash.deep == that.hash.deep
} 
Example 82
Source File: AvroToRowSpec.scala    From streamliner-examples   with Apache License 2.0 5 votes vote down vote up
package com.memsql.spark.examples.avro

import com.memsql.spark.connector.dataframe.JsonValue

import org.apache.avro.Schema
import org.apache.avro.generic.GenericData
import org.apache.spark.sql.Row
import test.util.Fixtures

import collection.JavaConversions._
import java.nio.ByteBuffer
import org.scalatest._

class AvroToRowSpec extends FlatSpec {
  "AvroToRow" should "create Spark SQL Rows from Avro objects" in {
    val parser: Schema.Parser = new Schema.Parser()
    val avroTestSchema: Schema = parser.parse(Fixtures.avroSchema)

    val record: GenericData.Record = new GenericData.Record(avroTestSchema)

    record.put("testBool", true)
    record.put("testDouble", 19.88)
    record.put("testFloat", 3.19f)
    record.put("testInt", 1123)
    record.put("testLong", 2147483648L)
    record.put("testNull", null)
    record.put("testString", "Conor")
    record.put("testUnion", 17)

    val row: Row = new AvroToRow().getRow(record)

    assert(row.getAs[Boolean](0))
    assert(row.getAs[Double](1) == 19.88)
    assert(row.getAs[Float](2) == 3.19f)
    assert(row.getAs[Int](3) == 1123)
    assert(row.getAs[Long](4) == 2147483648L)
    assert(row.getAs[Null](5) == null)
    assert(row.getAs[String](6) == "Conor")
    assert(row.getAs[String](7) == "17")
  }
} 
Example 83
Source File: ThriftRandomGenerator.scala    From streamliner-examples   with Apache License 2.0 5 votes vote down vote up
package com.memsql.spark.examples.thrift

import collection.JavaConversions._
import java.lang.reflect.Method
import java.nio.ByteBuffer
import org.apache.thrift.{TBase, TFieldIdEnum}
import org.apache.thrift.protocol.{TField, TType}
import org.apache.thrift.meta_data._

import scala.util.Random

object ThriftRandomGenerator {
  val random = new Random
  val MAX_RECURSION_LEVEL = 5

  def next[F <: TFieldIdEnum](c: Class[_], level: Int = 0): Any = {
    if (level > MAX_RECURSION_LEVEL) {
      return null
    }
    val className = c.getName
    try {
      val tBaseClass = c.asInstanceOf[Class[TBase[_ <: TBase[_, _], F]]]
      val instance = tBaseClass.newInstance()
      val metaDataMap: Map[_ <: TFieldIdEnum, FieldMetaData] = FieldMetaData.getStructMetaDataMap(tBaseClass).toMap
      metaDataMap.foreach({ case (field, fieldMetaData) =>
        val valueMetaData = fieldMetaData.valueMetaData
        val value = getValue(valueMetaData, level)
        instance.setFieldValue(instance.fieldForId(field.getThriftFieldId), value)
      })
      instance
    } catch {
      case e: ClassCastException => throw new IllegalArgumentException(s"Class $className is not a subclass of org.apache.thrift.TBase")
    }
  }

  def getValue(valueMetaData: FieldValueMetaData, level: Int): Any = {
    if (level > MAX_RECURSION_LEVEL) {
      return null
    }
    valueMetaData.`type` match {
      case TType.BOOL => random.nextBoolean
      case TType.BYTE => random.nextInt.toByte
      case TType.I16 => random.nextInt.toShort
      case TType.I32 => random.nextInt
      case TType.I64 => random.nextLong
      case TType.DOUBLE => random.nextInt(5) * 0.25
      case TType.ENUM => {
        val enumClass = valueMetaData.asInstanceOf[EnumMetaData].enumClass
        getEnumValue(enumClass)
      }
      case TType.STRING => {
        val length: Int = 5 + random.nextInt(5)
        val s = (1 to length).map(x => ('a'.toInt + random.nextInt(26)).toChar).mkString
        if (valueMetaData.isBinary) {
          ByteBuffer.wrap(s.getBytes)
        } else {
          s
        }
      }
      case TType.LIST => {
        val elemMetaData = valueMetaData.asInstanceOf[ListMetaData].elemMetaData
        val length: Int = 5 + random.nextInt(5)
        val ret: java.util.List[Any] = (1 to length).map(x => getValue(elemMetaData, level + 1))
        ret
      }
      case TType.SET => {
        val elemMetaData = valueMetaData.asInstanceOf[SetMetaData].elemMetaData
        val length: Int = 5 + random.nextInt(5)
        val ret: Set[Any] = (1 to length).map(x => getValue(elemMetaData, level + 1)).toSet
        val javaSet: java.util.Set[Any] = ret
        javaSet
      }
      case TType.MAP => {
        val mapMetaData = valueMetaData.asInstanceOf[MapMetaData]
        val keyMetaData = mapMetaData.keyMetaData
        val mapValueMetaData = mapMetaData.valueMetaData
        val length: Int = 5 + random.nextInt(5)
        val ret: Map[Any, Any] = (1 to length).map(_ => {
          val mapKey = getValue(keyMetaData, level + 1)
          val mapValue = getValue(mapValueMetaData, level + 1)
          mapKey -> mapValue
        }).toMap
        val javaMap: java.util.Map[Any, Any] = ret
        javaMap
      }
      case TType.STRUCT => {
        val structClass = valueMetaData.asInstanceOf[StructMetaData].structClass
        next(structClass, level = level + 1)
      }
      case _ => null
    }
  }

  def getEnumValue(enumType: Class[_]): Any = {
    val enumConstants = enumType.getEnumConstants
    enumConstants(random.nextInt(enumConstants.length))
  }
} 
Example 84
Source File: ThriftToRowSpec.scala    From streamliner-examples   with Apache License 2.0 5 votes vote down vote up
package com.memsql.spark.examples.thrift

import com.memsql.spark.connector.dataframe.JsonValue
import collection.JavaConversions._
import java.nio.ByteBuffer
import org.scalatest._

class ThriftToRowSpec extends FlatSpec {
  "ThriftToRow" should "create Spark SQL Rows from Thrift objects" in {
    val testClassInstance = new TestClass(
      true,
      42.toByte,
      128.toShort,
      1024,
      2048.toLong,
      2.5,
      "test1",
      ByteBuffer.wrap("test2".getBytes),
      mapAsJavaMap(Map("value" -> "test3")).asInstanceOf[java.util.Map[String, String]],
      List("test4"),
      Set("test5"),
      TestEnum.FIRST_VALUE,
      new SubClass("test6")
    )
    val thriftToRow = new ThriftToRow(classOf[TestClass])
    val row = thriftToRow.getRow(testClassInstance)
    assert(row.getAs[Boolean](0))
    assert(row.getAs[Byte](1) == 42.toByte)
    assert(row.getAs[Short](2) == 128.toShort)
    assert(row.getAs[Int](3) == 1024)
    assert(row.getAs[Long](4) == 2048)
    assert(row.getAs[Double](5) == 2.5)
    assert(row.getAs[String](6) == "test1")
    assert(row.getAs[ByteBuffer](7) == ByteBuffer.wrap("test2".getBytes))
    val mapValue = row.getAs[JsonValue](8)
    assert(mapValue.value == "{\"value\":\"test3\"}")
    val listValue = row.getAs[JsonValue](9)
    assert(listValue.value == "[\"test4\"]")
    val setValue = row.getAs[JsonValue](10)
    assert(setValue.value == "[\"test5\"]")
    assert(row.getAs[Int](11) == TestEnum.FIRST_VALUE.getValue)
    val subClassValue = row.getAs[JsonValue](12)
    assert(subClassValue.value == "{\"string_value\":\"test6\"}")
  }
} 
Example 85
Source File: SerializerFactory.scala    From spark-http-stream   with BSD 2-Clause "Simplified" License 5 votes vote down vote up
package org.apache.spark.sql.execution.streaming.http

import java.nio.ByteBuffer
import org.apache.spark.serializer.SerializerInstance
import org.apache.spark.serializer.DeserializationStream
import org.apache.spark.serializer.SerializationStream
import java.io.OutputStream
import java.io.InputStream
import scala.reflect.ClassTag
import com.fasterxml.jackson.databind.ObjectMapper
import org.apache.spark.SparkConf
import org.apache.spark.serializer.JavaSerializer
import org.apache.spark.serializer.KryoSerializer


object SerializerFactory {
	val DEFAULT = new SerializerFactory {
		override def getSerializerInstance(serializerName: String): SerializerInstance = {
			serializerName.toLowerCase() match {
				case "kryo" ⇒
					new KryoSerializer(new SparkConf()).newInstance();
				case "java" ⇒
					new JavaSerializer(new SparkConf()).newInstance();
				case _ ⇒ throw new InvalidSerializerNameException(serializerName);
			}
		}
	}
}

trait SerializerFactory {
	def getSerializerInstance(serializerName: String): SerializerInstance;
} 
Example 86
Source File: FileOps.scala    From sbt-aws-lambda   with Apache License 2.0 5 votes vote down vote up
package com.gilt.aws.lambda

import java.io.{ File, RandomAccessFile }
import java.nio.ByteBuffer

object FileOps {

  def fileToBuffer(file: File): ByteBuffer = {
    val buffer = ByteBuffer.allocate(file.length().toInt)
    val aFile = new RandomAccessFile(file, "r")
    val inChannel = aFile.getChannel()
    while (inChannel.read(buffer) > 0) {}
    inChannel.close()
    buffer.rewind()
    buffer
  }

} 
Example 87
Source File: package.scala    From aws-kinesis-scala   with Apache License 2.0 5 votes vote down vote up
package jp.co.bizreach

import java.nio.ByteBuffer

import com.amazonaws.services.kinesisfirehose.model.{
  PutRecordRequest => AWSPutRecordRequest,
  Record => AWSRecord,
  PutRecordResult => AWSPutRecordResult,
  PutRecordBatchRequest => AWSPutRecordBatchRequest,
  PutRecordBatchResult => AWSPutRecordBatchResult}

import scala.collection.JavaConverters._
import scala.language.implicitConversions

package object kinesisfirehose {

  private[this] implicit class JListConverters[A](list: java.util.List[A]) {
    def immutableSeq: Seq[A] = list.asScala.toSeq
  }

  case class PutRecordRequest(deliveryStreamName: String, record: Array[Byte])

  implicit def convertPutRecordRequest(request: PutRecordRequest): AWSPutRecordRequest = {
    val awsRequest = new AWSPutRecordRequest()
    awsRequest.setDeliveryStreamName(request.deliveryStreamName)
    awsRequest.setRecord(new AWSRecord().withData(ByteBuffer.wrap(request.record)))
    awsRequest
  }

  case class PutRecordResult(recordId: String)

  implicit def convertPutRecordResult(result: AWSPutRecordResult): PutRecordResult = {
    PutRecordResult(
      recordId = result.getRecordId
    )
  }

  case class PutRecordBatchRequest(deliveryStreamName: String, records: Seq[Array[Byte]])

  implicit def convertPutRecordBatchRequest(request: PutRecordBatchRequest): AWSPutRecordBatchRequest = {
    val awsRequest = new AWSPutRecordBatchRequest()
    awsRequest.setDeliveryStreamName(request.deliveryStreamName)
    awsRequest.setRecords(request.records.map { record =>
      new AWSRecord().withData(ByteBuffer.wrap(record))
    }.asJava)
    awsRequest
  }

  case class PutRecordBatchResult(failedPutCount: Int, records: Seq[PutRecordBatchResponseEntry])
  case class PutRecordBatchResponseEntry(recordId: String, errorCode: String, errorMessage: String)

  implicit def convertPutRecordBatchResult(result: AWSPutRecordBatchResult): PutRecordBatchResult = {
    PutRecordBatchResult(
      failedPutCount = result.getFailedPutCount,
      records = result.getRequestResponses.immutableSeq.map { record =>
        PutRecordBatchResponseEntry(
          recordId = record.getRecordId,
          errorCode = record.getErrorCode,
          errorMessage = record.getErrorMessage
        )
      }
    )
  }

} 
Example 88
Source File: Serdes.scala    From zio-kafka   with Apache License 2.0 5 votes vote down vote up
package zio.kafka.serde

import java.nio.ByteBuffer
import java.util.UUID

import org.apache.kafka.common.serialization.{ Serdes => KafkaSerdes }

private[zio] trait Serdes {
  lazy val long: Serde[Any, Long]             = Serde(KafkaSerdes.Long()).inmap(Long2long)(long2Long)
  lazy val int: Serde[Any, Int]               = Serde(KafkaSerdes.Integer()).inmap(Integer2int)(int2Integer)
  lazy val short: Serde[Any, Short]           = Serde(KafkaSerdes.Short()).inmap(Short2short)(short2Short)
  lazy val float: Serde[Any, Float]           = Serde(KafkaSerdes.Float()).inmap(Float2float)(float2Float)
  lazy val double: Serde[Any, Double]         = Serde(KafkaSerdes.Double()).inmap(Double2double)(double2Double)
  lazy val string: Serde[Any, String]         = Serde(KafkaSerdes.String())
  lazy val byteArray: Serde[Any, Array[Byte]] = Serde(KafkaSerdes.ByteArray())
  lazy val byteBuffer: Serde[Any, ByteBuffer] = Serde(KafkaSerdes.ByteBuffer())
  lazy val uuid: Serde[Any, UUID]             = Serde(KafkaSerdes.UUID())
} 
Example 89
Source File: Canvas.scala    From airframe   with Apache License 2.0 5 votes vote down vote up
package wvlet.airframe.canvas
import java.nio.ByteBuffer


  def newCanvas(size: Int): Canvas = {
    wrap(new Array[Byte](size))
  }

  def newOffHeapCanvas(size: Long): Canvas = {
    val m = defaultCanvasAllocator.allocate(size)
    new UnsafeCanvas(null, m.address, m.size, m)
  }

  def wrap(arr: Array[Byte]): Canvas                           = wrap(arr, 0, arr.length)
  def wrap(arr: Array[Byte], offset: Int, length: Int): Canvas = UnsafeCanvas.wrap(arr, offset, length)
  def wrap(buf: ByteBuffer): Canvas                            = UnsafeCanvas.wrap(buf)
} 
Example 90
Source File: L6-18Cassandra.scala    From prosparkstreaming   with Apache License 2.0 5 votes vote down vote up
package org.apress.prospark

import java.nio.charset.StandardCharsets
import org.apache.spark.SparkConf
import org.apache.spark.SparkContext
import org.apache.spark.streaming.Seconds
import org.apache.spark.streaming.StreamingContext
import org.json4s.DefaultFormats
import org.json4s.JField
import org.json4s.JsonAST.JObject
import org.json4s.jvalue2extractable
import org.json4s.jvalue2monadic
import org.json4s.native.JsonMethods.parse
import org.json4s.string2JsonInput
import org.apache.hadoop.conf.Configuration
import org.apache.hadoop.io.Text
import java.nio.ByteBuffer
import org.apache.cassandra.hadoop.ColumnFamilyOutputFormat
import org.apache.cassandra.hadoop.ConfigHelper
import org.apache.cassandra.thrift.ColumnOrSuperColumn
import org.apache.cassandra.thrift.Column
import org.apache.cassandra.utils.ByteBufferUtil
import org.apache.cassandra.thrift.Mutation
import java.util.Arrays

object CassandraSinkApp {

  def main(args: Array[String]) {
    if (args.length != 6) {
      System.err.println(
        "Usage: CassandraSinkApp <appname> <cassandraHost> <cassandraPort> <keyspace> <columnFamilyName> <columnName>")
      System.exit(1)
    }

    val Seq(appName, cassandraHost, cassandraPort, keyspace, columnFamilyName, columnName) = args.toSeq

    val conf = new SparkConf()
      .setAppName(appName)
      .setJars(SparkContext.jarOfClass(this.getClass).toSeq)

    val batchInterval = 10
    val windowSize = 20
    val slideInterval = 10

    val ssc = new StreamingContext(conf, Seconds(batchInterval))

    HttpUtils.createStream(ssc, url = "https://query.yahooapis.com/v1/public/yql?q=select%20*%20from%20yahoo.finance.quotes%20where%20symbol%20in%20(%22IBM,GOOG,MSFT,AAPL,FB,ORCL,YHOO,TWTR,LNKD,INTC%22)%0A%09%09&format=json&diagnostics=true&env=http%3A%2F%2Fdatatables.org%2Falltables.env",
      interval = batchInterval)
      .flatMap(rec => {
        implicit val formats = DefaultFormats
        val query = parse(rec) \ "query"
        ((query \ "results" \ "quote").children)
          .map(rec => ((rec \ "symbol").extract[String], (rec \ "LastTradePriceOnly").extract[String].toFloat))
      })
      .reduceByKeyAndWindow((x: Float, y: Float) => (x + y), Seconds(windowSize), Seconds(slideInterval))
      .foreachRDD(rdd => {
        val jobConf = new Configuration()
        ConfigHelper.setOutputRpcPort(jobConf, cassandraPort)
        ConfigHelper.setOutputInitialAddress(jobConf, cassandraHost)
        ConfigHelper.setOutputColumnFamily(jobConf, keyspace, columnFamilyName)
        ConfigHelper.setOutputPartitioner(jobConf, "Murmur3Partitioner")
        rdd.map(rec => {
          val c = new Column()
          c.setName(ByteBufferUtil.bytes(columnName))
          c.setValue(ByteBufferUtil.bytes(rec._2 / (windowSize / batchInterval)))
          c.setTimestamp(System.currentTimeMillis)
          val m = new Mutation()
          m.setColumn_or_supercolumn(new ColumnOrSuperColumn())
          m.column_or_supercolumn.setColumn(c)
          (ByteBufferUtil.bytes(rec._1), Arrays.asList(m))
        }).saveAsNewAPIHadoopFile(keyspace, classOf[ByteBuffer], classOf[List[Mutation]], classOf[ColumnFamilyOutputFormat], jobConf)
      })

    ssc.start()
    ssc.awaitTermination()
  }
} 
Example 91
Source File: ConnectedComponentsSpec.scala    From gemini   with GNU General Public License v3.0 5 votes vote down vote up
package tech.sourced.gemini

import java.nio.ByteBuffer

import org.slf4j.{Logger => Slf4jLogger}
import org.scalatest.{FlatSpec, Matchers}
import tech.sourced.gemini.util.Logger

import scala.collection.mutable.ListBuffer

class TestConnectedComponents(log: Slf4jLogger) extends ConnectedComponents(log) {
  def getHashtables(): List[Byte] = List(0, 1, 2)

  def intToHash(x: Byte): ByteBuffer = ByteBuffer.wrap(Array[Byte](x))

  // emulate database, restrictions:
  // - each hashtable must have all elements
  // - results must be sorted by (hash values, key)
  def getHashValues(hashtable: Byte): Iterable[FileHash] = {
    hashtable match {
      // bucket for a&b and d&3
      case 0 => List(
        FileHash("a", intToHash(1)),
        FileHash("b", intToHash(1)),
        FileHash("c", intToHash(2)),
        FileHash("d", intToHash(3)),
        FileHash("e", intToHash(3))
      )
      // bucket for b&c
      case 1 => List(
        FileHash("a", intToHash(1)),
        FileHash("b", intToHash(2)),
        FileHash("c", intToHash(2)),
        FileHash("d", intToHash(3)),
        FileHash("e", intToHash(4))
      )
      // no bucket
      case 2 => List(
        FileHash("a", intToHash(1)),
        FileHash("b", intToHash(2)),
        FileHash("c", intToHash(3)),
        FileHash("d", intToHash(4)),
        FileHash("e", intToHash(5))
      )
    }
  }
}

class ConnectedComponentsSpec extends FlatSpec
  with Matchers {

    val logger = Logger("gemini")
    val cc = new TestConnectedComponents(logger)

    "makeBuckets" should "correctly create buckets" in {
      cc.makeBuckets()._1 shouldEqual List[List[Int]](
        // buckets from hashtable 0
        List(0, 1),
        List(3, 4),
        // bucket from hashtable 1
        List(1, 2)
      )
    }

    "elementsToBuckets" should "create correct map" in {
      val (buckets, _) = cc.makeBuckets()
      cc.elementsToBuckets(buckets) shouldEqual Map[Int, List[Int]](
        0 -> List(0),
        1 -> List(0, 2),
        2 -> List(2),
        3 -> List(1),
        4 -> List(1)
      )
    }

  "findInBuckets" should "return connected components" in {
    val (buckets, _) = cc.makeBuckets()
    val elementToBuckets = cc.elementsToBuckets(buckets)
    cc.findInBuckets(buckets, elementToBuckets) shouldEqual Map[Int, Set[Int]](
      0 -> Set(1, 2, 0),
      1 -> Set(3, 4)
    )
  }
} 
Example 92
Source File: TestCompileApplicationInstance.scala    From milan   with Apache License 2.0 5 votes vote down vote up
package com.amazon.milan.tools

import java.io.{OutputStream, OutputStreamWriter}
import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets
import java.nio.file.Files

import com.amazon.milan.application.{Application, ApplicationConfiguration, ApplicationInstance}
import com.amazon.milan.lang._
import com.amazon.milan.testing.applications._
import com.amazon.milan.{Id, SemanticVersion}
import org.junit.Assert._
import org.junit.Test


object TestCompileApplicationInstance {

  case class Record(recordId: String, i: Int)

  class Provider extends ApplicationInstanceProvider {
    override def getApplicationInstance(params: List[(String, String)]): ApplicationInstance = {
      val input = Stream.of[Record]
      val graph = new StreamGraph(input)
      val config = new ApplicationConfiguration
      config.setListSource(input, Record("1", 1))

      val instanceId = params.find(_._1 == "instanceId").get._2
      val appId = params.find(_._1 == "appId").get._2

      new ApplicationInstance(
        instanceId,
        new Application(appId, graph, SemanticVersion.ZERO),
        config)
    }
  }

  class Compiler extends ApplicationInstanceCompiler {
    override def compile(applicationInstance: ApplicationInstance,
                         params: List[(String, String)],
                         output: OutputStream): Unit = {
      val writer = new OutputStreamWriter(output)
      val testParam = params.find(_._1 == "test").get._2
      writer.write(testParam)
      writer.write(applicationInstance.toJsonString)
      writer.close()
    }
  }

}

@Test
class TestCompileApplicationInstance {
  @Test
  def test_CompileApplicationInstance_Main_SendsProviderAndCompilerParameters(): Unit = {

    val tempFile = Files.createTempFile("TestCompileApplicationInstance", ".scala")
    Files.deleteIfExists(tempFile)

    val appId = Id.newId()
    val instanceId = Id.newId()
    val testValue = Id.newId()

    try {
      val args = Array(
        "--provider",
        "com.amazon.milan.tools.TestCompileApplicationInstance.Provider",
        "--compiler",
        "com.amazon.milan.tools.TestCompileApplicationInstance.Compiler",
        "--package",
        "generated",
        "--output",
        tempFile.toString,
        s"-PinstanceId=$instanceId",
        s"-PappId=$appId",
        s"-Ctest=$testValue"
      )
      CompileApplicationInstance.main(args)

      val fileContents = StandardCharsets.UTF_8.decode(ByteBuffer.wrap(Files.readAllBytes(tempFile))).toString
      assertTrue(fileContents.contains(appId))
      assertTrue(fileContents.contains(instanceId))
      assertTrue(fileContents.contains(testValue))
    }
    finally {
      Files.deleteIfExists(tempFile)
    }
  }
} 
Example 93
Source File: FlinkGenerator.scala    From milan   with Apache License 2.0 5 votes vote down vote up
package com.amazon.milan.compiler.flink.generator

import java.io.{ByteArrayOutputStream, OutputStream}
import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets
import java.nio.file.{Files, Path, StandardOpenOption}

import com.amazon.milan.application.{Application, ApplicationConfiguration, ApplicationInstance}
import com.amazon.milan.compiler.flink.internal.FlinkTypeEmitter
import com.amazon.milan.lang.StreamGraph
import com.amazon.milan.program.{Cycle, StreamExpression}
import com.amazon.milan.{Id, SemanticVersion}
import com.typesafe.scalalogging.Logger
import org.slf4j.LoggerFactory


case class GeneratorConfig(preventGenericTypeInformation: Boolean = false)


object FlinkGenerator {
  val default = new FlinkGenerator(GeneratorConfig())
}


class FlinkGenerator(classLoader: ClassLoader, generatorConfig: GeneratorConfig) {

  private val generatorTypeLifter = new FlinkTypeLifter(new FlinkTypeEmitter, this.generatorConfig.preventGenericTypeInformation)

  private val logger = Logger(LoggerFactory.getLogger(getClass))

  def this(generatorConfig: GeneratorConfig) {
    this(getClass.getClassLoader, generatorConfig)
  }

  def generateScala(graph: StreamGraph,
                    appConfig: ApplicationConfiguration,
                    packageName: String,
                    className: String): String = {
    val application = new Application(Id.newId(), graph, SemanticVersion.ZERO)
    val instance = new ApplicationInstance(Id.newId(), application, appConfig)
    this.generateScala(instance, packageName, className)
  }

  def generateScala(instance: ApplicationInstance,
                    outputPath: Path,
                    packageName: String,
                    className: String): Unit = {
    val scalaCode = this.generateScala(instance, packageName, className)
    val contents = scalaCode.getBytes(StandardCharsets.UTF_8)
    Files.write(outputPath, contents, StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING)
  }

  def generateScala(instance: ApplicationInstance,
                    packageName: String,
                    className: String): String = {
    val output = new ByteArrayOutputStream()
    this.generateScala(instance, output, packageName, className)

    output.flush()
    StandardCharsets.UTF_8.decode(ByteBuffer.wrap(output.toByteArray)).toString
  }

  def generateScala(instance: ApplicationInstance,
                    output: OutputStream,
                    packageName: String,
                    className: String): Unit = {
    val finalGraph = instance.application.graph.getDereferencedGraph
    finalGraph.typeCheckGraph()

    val outputs = new GeneratorOutputs(this.generatorTypeLifter)
    val context = GeneratorContext.createEmpty(instance.instanceDefinitionId, finalGraph, instance.config, outputs, this.generatorTypeLifter)

    // Ensure that every data stream is generated.
    finalGraph
      .getStreams
      .foreach(stream => this.ensureStreamIsGenerated(context, stream))

    // Close any cycles.
    finalGraph
      .getStreams
      .filter(_.isInstanceOf[Cycle])
      .map(_.asInstanceOf[Cycle])
      .foreach(context.closeCycle)

    // Add all sinks at the end.
    instance.config.dataSinks.foreach(sink => context.generateSink(sink))

    val generated = context.output.generateScala(packageName, className)
    output.write(generated.getBytes(StandardCharsets.UTF_8))
  }

  private def ensureStreamIsGenerated(context: GeneratorContext,
                                      stream: StreamExpression): Unit = {
    context.getOrGenerateDataStream(stream)
  }
} 
Example 94
Source File: FiloRowReader.scala    From filo   with Apache License 2.0 5 votes vote down vote up
package org.velvia.filo

import java.nio.ByteBuffer


class FastFiloRowReader(val parsers: Array[FiloVector[_]]) extends FiloRowReader {
  var rowNo: Int = -1
  def setRowNo(newRowNo: Int): Unit = { rowNo = newRowNo }

  def this(chunks: Array[ByteBuffer], classes: Array[Class[_]], emptyLen: Int = 0) =
    this(FiloVector.makeVectors(chunks, classes, emptyLen))

  final def notNull(columnNo: Int): Boolean = parsers(columnNo).isAvailable(rowNo)
  final def getBoolean(columnNo: Int): Boolean = parsers(columnNo).asInstanceOf[FiloVector[Boolean]](rowNo)
  final def getInt(columnNo: Int): Int = parsers(columnNo).asInstanceOf[FiloVector[Int]](rowNo)
  final def getLong(columnNo: Int): Long = parsers(columnNo).asInstanceOf[FiloVector[Long]](rowNo)
  final def getDouble(columnNo: Int): Double = parsers(columnNo).asInstanceOf[FiloVector[Double]](rowNo)
  final def getFloat(columnNo: Int): Float = parsers(columnNo).asInstanceOf[FiloVector[Float]](rowNo)
  final def getString(columnNo: Int): String = parsers(columnNo).asInstanceOf[FiloVector[String]](rowNo)
  override final def filoUTF8String(columnNo: Int): ZeroCopyUTF8String =
    parsers(columnNo).asInstanceOf[FiloVector[ZeroCopyUTF8String]](rowNo)
  final def getAny(columnNo: Int): Any = parsers(columnNo).boxed(rowNo)
}

// A RowReader that can safely be used in Seqs.  IE the rowNo is final and won't change.
case class SafeFiloRowReader(reader: FiloRowReader, rowNo: Int) extends FiloRowReader {
  val parsers = reader.parsers
  require(rowNo < parsers(0).length)
  def setRowNo(newRowNo: Int): Unit = {}

  final def notNull(columnNo: Int): Boolean = parsers(columnNo).isAvailable(rowNo)
  final def getBoolean(columnNo: Int): Boolean = parsers(columnNo).asInstanceOf[FiloVector[Boolean]](rowNo)
  final def getInt(columnNo: Int): Int = parsers(columnNo).asInstanceOf[FiloVector[Int]](rowNo)
  final def getLong(columnNo: Int): Long = parsers(columnNo).asInstanceOf[FiloVector[Long]](rowNo)
  final def getDouble(columnNo: Int): Double = parsers(columnNo).asInstanceOf[FiloVector[Double]](rowNo)
  final def getFloat(columnNo: Int): Float = parsers(columnNo).asInstanceOf[FiloVector[Float]](rowNo)
  final def getString(columnNo: Int): String = parsers(columnNo).asInstanceOf[FiloVector[String]](rowNo)
  override final def filoUTF8String(columnNo: Int): ZeroCopyUTF8String =
    parsers(columnNo).asInstanceOf[FiloVector[ZeroCopyUTF8String]](rowNo)
  final def getAny(columnNo: Int): Any = parsers(columnNo).boxed(rowNo)
} 
Example 95
Source File: RowToVectorBuilder.scala    From filo   with Apache License 2.0 5 votes vote down vote up
package org.velvia.filo

import java.nio.ByteBuffer
import scala.language.existentials
import scala.language.postfixOps
import scalaxy.loops._

import BuilderEncoder.{EncodingHint, AutoDetect}

case class VectorInfo(name: String, dataType: Class[_])

// To help matching against the ClassTag in the VectorBuilder
private object Classes {
  val Boolean = classOf[Boolean]
  val Byte = java.lang.Byte.TYPE
  val Short = java.lang.Short.TYPE
  val Int = java.lang.Integer.TYPE
  val Long = java.lang.Long.TYPE
  val Float = java.lang.Float.TYPE
  val Double = java.lang.Double.TYPE
  val String = classOf[String]
  val DateTime = classOf[org.joda.time.DateTime]
  val SqlTimestamp = classOf[java.sql.Timestamp]
  val UTF8 = classOf[ZeroCopyUTF8String]
}

object RowToVectorBuilder {
  
  def convertToBytes(hint: EncodingHint = AutoDetect): Map[String, ByteBuffer] = {
    val chunks = builders.map(_.toFiloBuffer(hint))
    schema.zip(chunks).map { case (VectorInfo(colName, _), bytes) => (colName, bytes) }.toMap
  }

  private def unsupportedInput(typ: Any) =
    throw new RuntimeException("Unsupported input type " + typ)
} 
Example 96
Source File: DiffEncoders.scala    From filo   with Apache License 2.0 5 votes vote down vote up
package org.velvia.filo.codecs

import com.google.flatbuffers.FlatBufferBuilder
import java.nio.{ByteBuffer, ByteOrder}
import org.joda.time.DateTime
import scala.collection.mutable.BitSet

import org.velvia.filo._
import org.velvia.filo.vector._


  def toPrimitiveVector[A: PrimitiveDataVectBuilder](data: Seq[A],
                                                     naMask: BitSet,
                                                     min: A,
                                                     max: A): ByteBuffer = {
    import DiffPrimitiveVector._

    val vectBuilder = implicitly[PrimitiveDataVectBuilder[A]]
    count += 1
    val fbb = new FlatBufferBuilder(getBuffer)
    val naOffset = populateNaMask(fbb, naMask, data.length)

    val ((dataOffset, dnbits), dsigned) = vectBuilder.buildDeltas(fbb, data, min, max)
    startDiffPrimitiveVector(fbb)
    addNaMask(fbb, naOffset)
    addLen(fbb, data.length)
    addData(fbb, dataOffset)
    addInfo(fbb, DataInfo.createDataInfo(fbb, dnbits, dsigned))
    addBase(fbb, vectBuilder.toLong(min))
    finishDiffPrimitiveVectorBuffer(fbb, endDiffPrimitiveVector(fbb))
    putHeaderAndGet(fbb, WireFormat.VECTORTYPE_DIFF, WireFormat.SUBTYPE_PRIMITIVE)
  }

  def toDateTimeVector(millis: LongVectorBuilder,
                       tz: IntVectorBuilder,
                       naMask: BitSet): ByteBuffer = {
    import DiffDateTimeVector._

    val intVectBuilder = AutoIntegralDVBuilders.IntDataVectBuilder
    val longVectBuilder = AutoIntegralDVBuilders.LongDataVectBuilder
    count += 1
    val fbb = new FlatBufferBuilder(getBuffer)
    val naOffset = populateNaMask(fbb, naMask, millis.length)

    val ((mOffset, mnbits), msigned) = longVectBuilder.buildDeltas(fbb, millis.data,
                                                               millis.min, millis.max)
    // Only build timezone vector if they are different.  Most DateTime's have same TZ
    val ((tOffset, tnbits), tsigned) = if (tz.min != tz.max) {
      intVectBuilder.buildDeltas(fbb, tz.data, tz.min, tz.max)
    } else {
      ((-1, -1), false)
    }

    startDiffDateTimeVector(fbb)
    addNaMask(fbb, naOffset)
    addVars(fbb, DDTVars.createDDTVars(fbb, millis.length, tz.min.toByte, millis.min))
    addMillisInfo(fbb, DataInfo.createDataInfo(fbb, mnbits, msigned))
    addMillis(fbb, mOffset)
    if (tOffset >= 0) {
      addTzInfo(fbb, DataInfo.createDataInfo(fbb, tnbits, tsigned))
      addTz(fbb, tOffset)
    }
    finishDiffDateTimeVectorBuffer(fbb, endDiffDateTimeVector(fbb))
    putHeaderAndGet(fbb, WireFormat.VECTORTYPE_DIFF, WireFormat.SUBTYPE_DATETIME)
  }
} 
Example 97
Source File: ConstEncoders.scala    From filo   with Apache License 2.0 5 votes vote down vote up
package org.velvia.filo.codecs

import com.google.flatbuffers.FlatBufferBuilder
import java.nio.{ByteBuffer, ByteOrder}
import scala.collection.mutable.BitSet

import org.velvia.filo._
import org.velvia.filo.vector._


  def toPrimitiveVector[A: PrimitiveDataVectBuilder](data: Seq[A],
                                                     naMask: BitSet,
                                                     min: A,
                                                     max: A): ByteBuffer = {
    import SimplePrimitiveVector._
    require(min == max)

    val vectBuilder = implicitly[PrimitiveDataVectBuilder[A]]
    count += 1
    val fbb = new FlatBufferBuilder(getBuffer)
    val naOffset = populateNaMask(fbb, naMask, data.length)
    val ((dataOffset, nbits), signed) = vectBuilder.build(fbb, Seq(min), min, max)
    startSimplePrimitiveVector(fbb)
    addNaMask(fbb, naOffset)
    addLen(fbb, data.length)
    addData(fbb, dataOffset)
    addInfo(fbb, DataInfo.createDataInfo(fbb, nbits, signed))
    finishSimplePrimitiveVectorBuffer(fbb, endSimplePrimitiveVector(fbb))
    putHeaderAndGet(fbb, WireFormat.VECTORTYPE_CONST, WireFormat.SUBTYPE_PRIMITIVE)
  }

  def toStringVector(str: String, len: Int, naMask: BitSet): ByteBuffer = {
    import ConstStringVector._

    count += 1
    val fbb = new FlatBufferBuilder(getBuffer)
    val naOffset = populateNaMask(fbb, naMask, len)
    val strOffset = fbb.createString(str)
    val offset = createConstStringVector(fbb, len, naOffset, strOffset)
    finishConstStringVectorBuffer(fbb, offset)
    putHeaderAndGet(fbb, WireFormat.VECTORTYPE_CONST, WireFormat.SUBTYPE_STRING)
  }
} 
Example 98
Source File: SimpleEncoders.scala    From filo   with Apache License 2.0 5 votes vote down vote up
package org.velvia.filo.codecs

import com.google.flatbuffers.FlatBufferBuilder
import java.nio.{ByteBuffer, ByteOrder}
import scala.collection.mutable.BitSet

import org.velvia.filo._
import org.velvia.filo.vector._


  def toPrimitiveVector[A: PrimitiveDataVectBuilder](data: Seq[A],
                                                     naMask: BitSet,
                                                     min: A,
                                                     max: A): ByteBuffer = {
    import SimplePrimitiveVector._

    val vectBuilder = implicitly[PrimitiveDataVectBuilder[A]]
    count += 1
    val fbb = new FlatBufferBuilder(getBuffer)
    val naOffset = populateNaMask(fbb, naMask, data.length)
    val ((dataOffset, nbits), signed) = vectBuilder.build(fbb, data, min, max)
    startSimplePrimitiveVector(fbb)
    addNaMask(fbb, naOffset)
    addLen(fbb, data.length)
    addData(fbb, dataOffset)
    addInfo(fbb, DataInfo.createDataInfo(fbb, nbits, signed))
    finishSimplePrimitiveVectorBuffer(fbb, endSimplePrimitiveVector(fbb))
    putHeaderAndGet(fbb, WireFormat.VECTORTYPE_SIMPLE, WireFormat.SUBTYPE_PRIMITIVE)
  }

  def toEmptyVector(len: Int): ByteBuffer = {
    val bb = ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN)
    bb.putInt(WireFormat.emptyVector(len))
    bb.position(0)
    bb
  }

  def toStringVector(data: Seq[String], naMask: BitSet): ByteBuffer = {
    val fbb = new FlatBufferBuilder(getBuffer)
    val naOffset = populateNaMask(fbb, naMask, data.length)
    val dataOffset = stringVect(fbb, data)
    val ssvOffset = SimpleStringVector.createSimpleStringVector(fbb, naOffset, dataOffset)
    SimpleStringVector.finishSimpleStringVectorBuffer(fbb, ssvOffset)
    putHeaderAndGet(fbb, WireFormat.VECTORTYPE_SIMPLE, WireFormat.SUBTYPE_STRING)
  }
} 
Example 99
Source File: DictEncodingEncoders.scala    From filo   with Apache License 2.0 5 votes vote down vote up
package org.velvia.filo.codecs

import com.google.flatbuffers.FlatBufferBuilder
import java.nio.ByteBuffer
import java.util.HashMap
import scala.collection.mutable.{ArrayBuffer, BitSet}
import scala.language.postfixOps
import scalaxy.loops._

import org.velvia.filo._
import org.velvia.filo.vector._


object DictEncodingEncoders extends ThreadLocalBuffers {
  import Utils._

  var count = 0

  // Note: This is a way to avoid storing null and dealing with NPEs for NA values
  val NaString = ""

  def toStringVector(data: Seq[String], naMask: BitSet, stringSet: collection.Set[String]): ByteBuffer = {
    import DictStringVector._

    count += 1
    val builder = AutoIntegralDVBuilders.IntDataVectBuilder

    // Convert the set of strings to an encoding
    val uniques = stringSet.toSeq
    // NOTE: sorry but java's HashMap is just much faster (for the next step)
    // This used to be `uniques.zipWithIndex.toMap`
    val strToCode = new HashMap[String, Int]()
    for { i <- 0 until uniques.length optimized } {
      strToCode.put(uniques(i), i)
    }

    // Encode each string to the code per the map above
    // Again we could have used data.zipWithIndex.map(....) but this is much faster.
    val codes = ArrayBuffer.fill(data.length)(0)
    for { i <- 0 until data.length optimized } {
      if (!naMask(i)) codes(i) = strToCode.get(data(i)) + 1
    }

    val fbb = new FlatBufferBuilder(getBuffer)
    val ((dataOffset, nbits), signed) = builder.build(fbb, codes, 0, stringSet.size + 1)
    val dictVect = stringVect(fbb, Seq(NaString) ++ uniques)
    startDictStringVector(fbb)
    addDictionary(fbb, dictVect)
    addLen(fbb, data.length)
    addCodes(fbb, dataOffset)
    addInfo(fbb, DataInfo.createDataInfo(fbb, nbits, signed))
    finishDictStringVectorBuffer(fbb, endDictStringVector(fbb))
    putHeaderAndGet(fbb, WireFormat.VECTORTYPE_DICT, WireFormat.SUBTYPE_STRING)
  }
} 
Example 100
Source File: MemberMetadata.scala    From CMAK   with Apache License 2.0 5 votes vote down vote up
package kafka.manager.utils.one10
import java.nio.ByteBuffer

import org.apache.kafka.clients.admin.{ConsumerGroupDescription, MemberDescription}
import org.apache.kafka.clients.consumer.internals.ConsumerProtocol
import org.apache.kafka.common.requests.DescribeGroupsResponse
import org.apache.kafka.common.utils.Utils

object MemberMetadata {
  import collection.JavaConverters._
  def from(groupId: String, groupSummary: ConsumerGroupDescription, memberSummary: MemberDescription) : MemberMetadata = {
    val assignment = memberSummary.assignment().topicPartitions().asScala.map(tp => tp.topic() -> tp.partition()).toSet
    MemberMetadata(
      memberSummary.consumerId()
      , groupId
      , memberSummary.clientId
      , memberSummary.host()
      , "(n/a on backfill)"
      , List.empty
      , assignment
    )
  }
}



case class MemberMetadata(memberId: String,
                          groupId: String,
                          clientId: String,
                          clientHost: String,
                          protocolType: String,
                          supportedProtocols: List[(String, Set[String])],
                          assignment: Set[(String, Int)]
                    ) {

  def protocols = supportedProtocols.map(_._1).toSet

  def metadata(protocol: String): Set[String] = {
    supportedProtocols.find(_._1 == protocol) match {
      case Some((_, metadata)) => metadata
      case None =>
        throw new IllegalArgumentException("Member does not support protocol")
    }
  }

  override def toString: String = {
    "MemberMetadata(" +
      s"memberId=$memberId, " +
      s"clientId=$clientId, " +
      s"clientHost=$clientHost, " +
      s"supportedProtocols=${supportedProtocols.map(_._1)}, " +
      ")"
  }

} 
Example 101
Source File: MemberMetadata.scala    From CMAK   with Apache License 2.0 5 votes vote down vote up
package kafka.manager.utils.two40
import java.nio.ByteBuffer

import org.apache.kafka.clients.admin.{ConsumerGroupDescription, MemberDescription}
import org.apache.kafka.clients.consumer.internals.ConsumerProtocol
import org.apache.kafka.common.requests.DescribeGroupsResponse
import org.apache.kafka.common.utils.Utils

object MemberMetadata {
  import collection.JavaConverters._
  def from(groupId: String, groupSummary: ConsumerGroupDescription, memberSummary: MemberDescription) : MemberMetadata = {
    val assignment = memberSummary.assignment().topicPartitions().asScala.map(tp => tp.topic() -> tp.partition()).toSet
    MemberMetadata(
      memberSummary.consumerId()
      , groupId
      , None
      , memberSummary.clientId
      , memberSummary.host()
      , "(n/a on backfill)"
      , List.empty
      , assignment
    )
  }
}



case class MemberMetadata(memberId: String,
                          groupId: String,
                          groupInstanceId: Option[String],
                          clientId: String,
                          clientHost: String,
                          protocolType: String,
                          supportedProtocols: List[(String, Set[String])],
                          assignment: Set[(String, Int)]
                    ) {

  def protocols = supportedProtocols.map(_._1).toSet

  def metadata(protocol: String): Set[String] = {
    supportedProtocols.find(_._1 == protocol) match {
      case Some((_, metadata)) => metadata
      case None =>
        throw new IllegalArgumentException("Member does not support protocol")
    }
  }

  override def toString: String = {
    "MemberMetadata(" +
      s"memberId=$memberId, " +
      s"groupInstanceId=$groupInstanceId, " +
      s"clientId=$clientId, " +
      s"clientHost=$clientHost, " +
      s"supportedProtocols=${supportedProtocols.map(_._1)}, " +
      ")"
  }

} 
Example 102
Source File: MessageWriter.scala    From lsp4s   with Apache License 2.0 5 votes vote down vote up
package scala.meta.jsonrpc

import java.io.ByteArrayOutputStream
import java.io.OutputStream
import java.io.OutputStreamWriter
import java.io.PrintWriter
import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets
import scala.concurrent.Future
import io.circe.syntax._
import monix.execution.Ack
import monix.reactive.Observer
import scribe.LoggerSupport


  def write(msg: Message): Future[Ack] = lock.synchronized {
    baos.reset()
    val json = msg.asJson
    val protocol = BaseProtocolMessage.fromJson(json)
    logger.trace(s" --> $json")
    val byteBuffer = MessageWriter.write(protocol, baos, headerOut)
    out.onNext(byteBuffer)
  }
}

object MessageWriter {

  def headerWriter(out: OutputStream): PrintWriter = {
    new PrintWriter(new OutputStreamWriter(out, StandardCharsets.US_ASCII))
  }

  def write(message: BaseProtocolMessage): ByteBuffer = {
    val out = new ByteArrayOutputStream()
    val header = headerWriter(out)
    write(message, out, header)
  }

  def write(
      message: BaseProtocolMessage,
      out: ByteArrayOutputStream,
      headerOut: PrintWriter
  ): ByteBuffer = {
    message.header.foreach {
      case (key, value) =>
        headerOut.write(key)
        headerOut.write(": ")
        headerOut.write(value)
        headerOut.write("\r\n")
    }
    headerOut.write("\r\n")
    out.write(message.content)
    out.flush()
    val buffer = ByteBuffer.wrap(out.toByteArray, 0, out.size())
    buffer
  }
} 
Example 103
Source File: MonixEnrichments.scala    From lsp4s   with Apache License 2.0 5 votes vote down vote up
package scala.meta.jsonrpc

import java.io.IOException
import java.io.OutputStream
import java.nio.ByteBuffer
import monix.execution.Ack
import monix.execution.Cancelable
import monix.execution.Scheduler
import monix.reactive.Observable
import monix.reactive.Observer
import scribe.LoggerSupport

object MonixEnrichments {

  
  class ObservableCurrentValue[+A](obs: Observable[A])(implicit s: Scheduler)
      extends (() => A)
      with Cancelable {
    private var value: Any = _
    private val cancelable = obs.foreach(newValue => value = newValue)
    override def apply(): A = {
      if (value == null) {
        throw new NoSuchElementException(
          "Reading from empty Observable, consider using MulticastStrategy.behavior(initialValue)"
        )
      } else {
        value.asInstanceOf[A]
      }
    }
    override def cancel(): Unit = cancelable.cancel()
  }

  implicit class XtensionObservable[A](val obs: Observable[A]) extends AnyVal {

    def focus[B: cats.Eq](f: A => B): Observable[B] =
      obs.distinctUntilChangedByKey(f).map(f)

    def toFunction0()(implicit s: Scheduler): () => A =
      toObservableCurrentValue()

    def toObservableCurrentValue()(
        implicit s: Scheduler
    ): ObservableCurrentValue[A] =
      new ObservableCurrentValue[A](obs)
  }

  implicit class XtensionObserverCompanion[A](val `_`: Observer.type)
      extends AnyVal {
    def fromOutputStream(
        out: OutputStream,
        logger: LoggerSupport
    ): Observer.Sync[ByteBuffer] = {
      new Observer.Sync[ByteBuffer] {
        private[this] var isClosed: Boolean = false
        override def onNext(elem: ByteBuffer): Ack = {
          if (isClosed) Ack.Stop
          else {
            try {
              while (elem.hasRemaining) out.write(elem.get())
              out.flush()
              Ack.Continue
            } catch {
              case _: IOException =>
                logger.error("OutputStream closed!")
                isClosed = true
                Ack.Stop
            }
          }
        }
        override def onError(ex: Throwable): Unit = ()
        override def onComplete(): Unit = out.close()
      }
    }
  }

} 
Example 104
Source File: LanguageClient.scala    From lsp4s   with Apache License 2.0 5 votes vote down vote up
package scala.meta.jsonrpc

import cats.syntax.either._
import io.circe.Decoder
import io.circe.Encoder
import io.circe.syntax._
import java.io.OutputStream
import java.nio.ByteBuffer
import monix.eval.Callback
import monix.eval.Task
import monix.execution.Ack
import monix.execution.Cancelable
import monix.execution.atomic.Atomic
import monix.execution.atomic.AtomicInt
import monix.reactive.Observer
import scala.collection.concurrent.TrieMap
import scala.concurrent.Future
import scala.concurrent.duration.Duration
import MonixEnrichments._
import scribe.LoggerSupport

class LanguageClient(out: Observer[ByteBuffer], logger: LoggerSupport)
    extends JsonRpcClient {
  def this(out: OutputStream, logger: LoggerSupport) =
    this(Observer.fromOutputStream(out, logger), logger)
  private val writer = new MessageWriter(out, logger)
  private val counter: AtomicInt = Atomic(1)
  private val activeServerRequests =
    TrieMap.empty[RequestId, Callback[Response]]
  def notify[A: Encoder](method: String, notification: A): Future[Ack] =
    writer.write(Notification(method, Some(notification.asJson)))
  def serverRespond(response: Response): Future[Ack] = response match {
    case Response.Empty => Ack.Continue
    case x: Response.Success => writer.write(x)
    case x: Response.Error =>
      logger.error(s"Response error: $x")
      writer.write(x)
  }
  def clientRespond(response: Response): Unit =
    for {
      id <- response match {
        case Response.Empty => None
        case Response.Success(_, requestId) => Some(requestId)
        case Response.Error(_, requestId) => Some(requestId)
      }
      callback <- activeServerRequests.get(id).orElse {
        logger.error(s"Response to unknown request: $response")
        None
      }
    } {
      activeServerRequests.remove(id)
      callback.onSuccess(response)
    }

  def request[A: Encoder, B: Decoder](
      method: String,
      request: A
  ): Task[Either[Response.Error, B]] = {
    val nextId = RequestId(counter.incrementAndGet())
    val response = Task.create[Response] { (out, cb) =>
      val scheduled = out.scheduleOnce(Duration(0, "s")) {
        val json = Request(method, Some(request.asJson), nextId)
        activeServerRequests.put(nextId, cb)
        writer.write(json)
      }
      Cancelable { () =>
        scheduled.cancel()
        this.notify("$/cancelRequest", CancelParams(nextId.value))
      }
    }
    response.map {
      case Response.Empty =>
        Left(
          Response.invalidParams(
            s"Got empty response for request $request",
            nextId
          )
        )
      case err: Response.Error =>
        Left(err)
      case Response.Success(result, _) =>
        result.as[B].leftMap { err =>
          Response.invalidParams(err.toString, nextId)
        }
    }
  }
}

object LanguageClient {
  def fromOutputStream(out: OutputStream, logger: LoggerSupport) =
    new LanguageClient(Observer.fromOutputStream(out, logger), logger)
} 
Example 105
Source File: BaseProtocolMessage.scala    From lsp4s   with Apache License 2.0 5 votes vote down vote up
package scala.meta.jsonrpc

import java.io.InputStream
import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets
import java.util
import io.circe.Json
import io.circe.syntax._
import monix.reactive.Observable
import scribe.LoggerSupport

final class BaseProtocolMessage(
    val header: Map[String, String],
    val content: Array[Byte]
) {

  override def equals(obj: scala.Any): Boolean =
    this.eq(obj.asInstanceOf[Object]) || {
      obj match {
        case m: BaseProtocolMessage =>
          header.equals(m.header) &&
            util.Arrays.equals(content, m.content)
      }
    }

  override def toString: String = {
    val bytes = MessageWriter.write(this)
    StandardCharsets.UTF_8.decode(bytes).toString
  }
}

object BaseProtocolMessage {
  val ContentLen = "Content-Length"

  def apply(msg: Message): BaseProtocolMessage =
    fromJson(msg.asJson)
  def fromJson(json: Json): BaseProtocolMessage =
    fromBytes(json.noSpaces.getBytes(StandardCharsets.UTF_8))
  def fromBytes(bytes: Array[Byte]): BaseProtocolMessage =
    new BaseProtocolMessage(
      Map("Content-Length" -> bytes.length.toString),
      bytes
    )

  def fromInputStream(
      in: InputStream,
      logger: LoggerSupport
  ): Observable[BaseProtocolMessage] =
    fromBytes(Observable.fromInputStream(in), logger)

  def fromBytes(
      in: Observable[Array[Byte]],
      logger: LoggerSupport
  ): Observable[BaseProtocolMessage] =
    fromByteBuffers(in.map(ByteBuffer.wrap), logger)

  def fromByteBuffers(
      in: Observable[ByteBuffer],
      logger: LoggerSupport
  ): Observable[BaseProtocolMessage] =
    in.executeAsync.liftByOperator(new BaseProtocolMessageParser(logger))
} 
Example 106
Source File: BosonExtractor.scala    From boson   with Apache License 2.0 5 votes vote down vote up
package io.zink.boson.impl

import java.nio.ByteBuffer

import io.zink.boson.Boson
import io.zink.boson.bson.bsonImpl.BosonImpl
import io.zink.boson.bson.bsonPath._
import shapeless.TypeCase

import scala.concurrent.ExecutionContext.Implicits.global
import scala.concurrent.Future


class BosonExtractor[T](expression: String, extractFunction: T => Unit)(implicit tp: Option[TypeCase[T]]) extends Boson {

  private val interpreter: Interpreter[T] = new Interpreter[T](expression, fExt = Some(extractFunction))(tp, None)

  private def runInterpreter(bsonEncoded: Either[Array[Byte], String]): Unit = {
    interpreter.run(bsonEncoded)
  }

  override def go(bsonByteEncoding: Array[Byte]): Future[Array[Byte]] = {
    val future: Future[Array[Byte]] = Future {
      runInterpreter(Left(bsonByteEncoding))
      bsonByteEncoding
    }
    future
  }

  override def go(bsonByteEncoding: String): Future[String] = {
    val future: Future[String] =
      Future {
        runInterpreter(Right(bsonByteEncoding))
        bsonByteEncoding
      }
    future
  }

//  override def fuse(boson: Boson): Boson = new BosonFuse(this, boson)

} 
Example 107
Source File: Interpreter.scala    From hammock   with MIT License 5 votes vote down vote up
package hammock
package js

import cats._
import cats.effect.{Async, ContextShift}
import cats.syntax.applicative._
import cats.syntax.flatMap._
import cats.syntax.functor._
import cats.syntax.show._
import org.scalajs.dom.ext.Ajax
import org.scalajs.dom.ext.Ajax.InputData
import java.nio.ByteBuffer

object Interpreter {

  def apply[F[_]](implicit F: InterpTrans[F]): InterpTrans[F] = F

  implicit def instance[F[_]: Async: ContextShift]: InterpTrans[F] = new InterpTrans[F] {
    def trans: HttpF ~> F = {

      def doReq(reqF: HttpF[HttpResponse]): F[HttpResponse] = {
        val timeout = 0
        val headers = reqF.req.headers
        val data: InputData = reqF.req.entity.fold(InputData.str2ajax(""))(
          _.cata(
            string => InputData.str2ajax(string.content),
            bytes => InputData.byteBuffer2ajax(ByteBuffer.wrap(bytes.content)),
            Function.const(InputData.str2ajax("")))
        )
        val method = toMethod(reqF)

        for {
          response <- Async.fromFuture(
            Async[F].delay(Ajax(method.name, reqF.req.uri.show, data, timeout, headers, false, "")))
          responseHeaders <- parseHeaders(response.getAllResponseHeaders)
          status = Status.get(response.status)
          body   = response.responseText
        } yield HttpResponse(status, responseHeaders, Entity.StringEntity(body))
      }

      def toMethod(reqF: HttpF[HttpResponse]): Method = reqF match {
        case Options(_) => Method.OPTIONS
        case Get(_)     => Method.GET
        case Head(_)    => Method.HEAD
        case Post(_)    => Method.POST
        case Put(_)     => Method.PUT
        case Delete(_)  => Method.DELETE
        case Trace(_)   => Method.TRACE
        case Patch(_)   => Method.PATCH
      }

      def parseHeaders(str: String): F[Map[String, String]] = str match {
        case null => Map.empty[String, String].pure[F]
        case string =>
          Async[F].delay(
            string
              .split("\r\n")
              .map({ line =>
                val splitted = line.split(": ")
                (splitted.head, splitted.tail.mkString("").trim)
              })
              .toMap)
      }

      λ[HttpF ~> F] {
        case req @ (Options(_) | Get(_) | Head(_) | Post(_) | Put(_) | Delete(_) | Trace(_) | Patch(_)) => doReq(req)
      }
    }
  }
} 
Example 108
Source File: MyNettyBlockRpcServer.scala    From OAP   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.network.netty

import java.nio.ByteBuffer

import scala.language.existentials

import org.apache.spark.SparkEnv
import org.apache.spark.internal.Logging
import org.apache.spark.network.BlockDataManager
import org.apache.spark.network.client.{RpcResponseCallback, StreamCallbackWithID, TransportClient}
import org.apache.spark.network.server.{OneForOneStreamManager, RpcHandler, StreamManager}
import org.apache.spark.network.shuffle.protocol._
import org.apache.spark.serializer.Serializer
import org.apache.spark.shuffle.remote.{HadoopFileSegmentManagedBuffer, MessageForHadoopManagedBuffers, RemoteShuffleManager}
import org.apache.spark.shuffle.sort.SortShuffleManager
import org.apache.spark.storage.{BlockId, ShuffleBlockId}


class MyNettyBlockRpcServer(
    appId: String,
    serializer: Serializer,
    blockManager: BlockDataManager)
  extends RpcHandler with Logging {

  private val streamManager = new OneForOneStreamManager()

  override def receive(
      client: TransportClient,
      rpcMessage: ByteBuffer,
      responseContext: RpcResponseCallback): Unit = {
    val message = BlockTransferMessage.Decoder.fromByteBuffer(rpcMessage)
    logTrace(s"Received request: $message")

    message match {
      case openBlocks: OpenBlocks =>
        val blocksNum = openBlocks.blockIds.length
        val isShuffleRequest = (blocksNum > 0) &&
          BlockId.apply(openBlocks.blockIds(0)).isInstanceOf[ShuffleBlockId] &&
          (SparkEnv.get.conf.get("spark.shuffle.manager", classOf[SortShuffleManager].getName)
            == classOf[RemoteShuffleManager].getName)
        if (isShuffleRequest) {
          val blockIdAndManagedBufferPair =
            openBlocks.blockIds.map(block => (block, blockManager.getHostLocalShuffleData(
              BlockId.apply(block), Array.empty).asInstanceOf[HadoopFileSegmentManagedBuffer]))
          responseContext.onSuccess(new MessageForHadoopManagedBuffers(
            blockIdAndManagedBufferPair).toByteBuffer.nioBuffer())
        } else {
          // This customized Netty RPC server is only served for RemoteShuffle requests,
          // Other RPC messages or data chunks transferring should go through
          // NettyBlockTransferService' NettyBlockRpcServer
          throw new UnsupportedOperationException("MyNettyBlockRpcServer only serves remote" +
            " shuffle requests for OpenBlocks")
        }

      case uploadBlock: UploadBlock =>
        throw new UnsupportedOperationException("MyNettyBlockRpcServer doesn't serve UploadBlock")
    }
  }

  override def receiveStream(
      client: TransportClient,
      messageHeader: ByteBuffer,
      responseContext: RpcResponseCallback): StreamCallbackWithID = {
    throw new UnsupportedOperationException("MyNettyBlockRpcServer doesn't support receiving" +
      " stream")
  }

  override def getStreamManager(): StreamManager = streamManager
} 
Example 109
Source File: Decode.scala    From scala-stellar-sdk   with Apache License 2.0 5 votes vote down vote up
package stellar.sdk.model.xdr

import java.io.EOFException
import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets
import java.time.Instant

import cats.Eval
import cats.data.{IndexedStateT, State}
import com.typesafe.scalalogging.LazyLogging

import scala.util.Try

trait Decode extends LazyLogging {

  private def decode[T](bs: Seq[Byte], len: Int)(decoder: Seq[Byte] => T): (Seq[Byte], T) = {
    if (bs.length < len) throw new EOFException("Insufficient data remains to parse.")
    val t = decoder(bs.take(len))
    logger.trace(s"Dropping {} to make {}", len, t)
    bs.drop(len) -> t
  }

  val int: State[Seq[Byte], Int] = State[Seq[Byte], Int] { bs =>
    decode(bs, 4) { in => ByteBuffer.wrap(in.toArray).getInt }
  }

  val long: State[Seq[Byte], Long] = State[Seq[Byte], Long] { bs =>
    decode(bs, 8) { in => ByteBuffer.wrap(in.toArray).getLong }
  }

  val instant: State[Seq[Byte], Instant] = long.map(Instant.ofEpochSecond)

  val bool: State[Seq[Byte], Boolean] = int.map(_ == 1)

  def bytes(len: Int): State[Seq[Byte], Seq[Byte]] = State[Seq[Byte], Seq[Byte]] { bs =>
    decode(bs, len) { _.take(len) }
  }

  val bytes: State[Seq[Byte], Seq[Byte]] = for {
    len <- int
    bs <- bytes(len)
  } yield bs

  def padded(multipleOf: Int = 4): State[Seq[Byte], Seq[Byte]] = for {
    len <- int
    bs <- bytes(len)
    _ <- bytes((multipleOf - (len % multipleOf)) % multipleOf)
  } yield bs

  val string: State[Seq[Byte], String] = padded().map(_.toArray).map(new String(_, StandardCharsets.UTF_8))

  def switch[T](zero: State[Seq[Byte], T], others: State[Seq[Byte], T]*): IndexedStateT[Eval, Seq[Byte], Seq[Byte], T] = int.flatMap {
    case 0 => zero
    case n =>  Try(others(n - 1)).getOrElse {
      throw new IllegalArgumentException(s"No parser defined for discriminant $n")
    }
  }

  // TODO (jem) - All switches should use this instead and Discriminators should be held in the parent (switcher not switchee).
  def switchInt[T](zero: State[Seq[Byte], T], others: State[Seq[Byte], T]*): State[Seq[Byte], (T, Int)] = int.flatMap {
    case 0 => zero.map(_ -> 0)
    case n => Try(others(n - 1).map(_ -> n)).getOrElse {
      throw new IllegalArgumentException(s"No parser defined for discriminant $n")
    }
  }

  def opt[T](parseT: State[Seq[Byte], T]): State[Seq[Byte], Option[T]] = bool.flatMap {
    case true => parseT.map(Some(_))
    case false => State.pure(None)
  }

  def arr[T](parseT: State[Seq[Byte], T]): State[Seq[Byte], Seq[T]] = int.flatMap(seq(_, parseT))

  // $COVERAGE-OFF$
  // For debugging XDR only.
  def log[T](t: T): State[Seq[Byte], Unit] = State[Seq[Byte], Unit] { bs =>
    logger.debug("{}\n", t)
    bs -> ()
  }
  // $COVERAGE-ON$

  def seq[T](qty: Int, parseT: State[Seq[Byte], T]): State[Seq[Byte], Seq[T]] = {
    (0 until qty).foldLeft(State.pure[Seq[Byte], Seq[T]](Seq.empty[T])) { case (state, _) =>
      for {
        ts <- state
        t <- parseT
      } yield ts :+ t
    }
  }

  def drop[T](parse: State[Seq[Byte], _])(t: T): State[Seq[Byte], T] = for {
    _ <- parse
  } yield t

  def widen[A, W, O <: W](s: State[A, O]): State[A, W] = s.map(w => w: W)
} 
Example 110
Source File: NMClient.scala    From incubator-retired-gearpump   with Apache License 2.0 5 votes vote down vote up
package org.apache.gearpump.experiments.yarn.glue

import java.nio.ByteBuffer

import akka.actor.ActorRef
import com.typesafe.config.Config
import org.apache.gearpump.experiments.yarn.appmaster.YarnAppMaster.ContainerStarted
import org.apache.gearpump.experiments.yarn.glue.Records._
import org.apache.gearpump.util.LogUtil
import org.apache.hadoop.yarn.api.records.{ContainerId => YarnContainerId, ContainerStatus => YarnContainerStatus}
import org.apache.hadoop.yarn.client.api.async.NMClientAsync
import org.apache.hadoop.yarn.client.api.async.impl.NMClientAsyncImpl

class NMClient(yarnConf: YarnConfig, config: Config) extends NMClientAsync.CallbackHandler {

  private val LOG = LogUtil.getLogger(getClass)

  private var reportTo: ActorRef = null
  private var client: NMClientAsyncImpl = null

  def start(reportTo: ActorRef): Unit = {
    LOG.info("Starting Node Manager Client NMClient...")
    this.reportTo = reportTo
    client = new NMClientAsyncImpl(this)
    client.init(yarnConf.conf)
    client.start()
  }

  private[glue]
  override def onContainerStarted(
      containerId: YarnContainerId, allServiceResponse: java.util.Map[String, ByteBuffer]) {
    LOG.info(s"Container started : $containerId, " + allServiceResponse)
    reportTo ! ContainerStarted(containerId)
  }

  private[glue]
  override def onContainerStatusReceived(
      containerId: YarnContainerId, containerStatus: YarnContainerStatus) {
    LOG.info(s"Container status received : $containerId, status $containerStatus")
  }

  private[glue]
  override def onContainerStopped(containerId: YarnContainerId) {
    LOG.error(s"Container stopped : $containerId")
  }

  private[glue]
  override def onGetContainerStatusError(containerId: YarnContainerId, throwable: Throwable) {
    LOG.error(s"Container exception : $containerId", throwable)
  }

  private[glue]
  override def onStartContainerError(containerId: YarnContainerId, throwable: Throwable) {
    LOG.error(s"Container exception : $containerId", throwable)
  }

  private[glue]
  override def onStopContainerError(containerId: YarnContainerId, throwable: Throwable) {
    LOG.error(s"Container exception : $containerId", throwable)
  }

  def launchCommand(
      container: Container, command: String, packagePath: String, configPath: String): Unit = {
    LOG.info(s"Launching command : $command on container" +
      s":  ${container.getId}, host ip : ${container.getNodeId.getHost}")
    val context = ContainerLaunchContext(yarnConf.conf, command, packagePath, configPath)
    client.startContainerAsync(container, context)
  }

  def stopContainer(containerId: ContainerId, nodeId: NodeId): Unit = {
    LOG.info(s"Stop container ${containerId.toString} on node: ${nodeId.toString} ")
    client.stopContainerAsync(containerId, nodeId)
  }

  def stop(): Unit = {
    LOG.info(s"Shutdown NMClient")
    client.stop()
  }
} 
Example 111
Source File: ContainerLaunchContext.scala    From incubator-retired-gearpump   with Apache License 2.0 5 votes vote down vote up
package org.apache.gearpump.experiments.yarn.glue

import java.io.File
import java.nio.ByteBuffer
import scala.collection.JavaConverters._

import org.apache.hadoop.fs.{FileSystem => YarnFileSystem, Path}
import org.apache.hadoop.io.DataOutputBuffer
import org.apache.hadoop.mapreduce.security.TokenCache
import org.apache.hadoop.security.UserGroupInformation
import org.apache.hadoop.yarn.api.ApplicationConstants.Environment
import org.apache.hadoop.yarn.api.records._
import org.apache.hadoop.yarn.conf.YarnConfiguration
import org.apache.hadoop.yarn.util.ConverterUtils
import org.slf4j.Logger

import org.apache.gearpump.util.LogUtil

private[glue]
object ContainerLaunchContext {
  private val LOG: Logger = LogUtil.getLogger(getClass)

  def apply(yarnConf: YarnConfiguration, command: String, packagePath: String, configPath: String)
    : ContainerLaunchContext = {
    val context = Records.newRecord(classOf[ContainerLaunchContext])
    context.setCommands(Seq(command).asJava)
    context.setEnvironment(getAppEnv(yarnConf).asJava)
    context.setTokens(getToken(yarnConf, packagePath, configPath))
    context.setLocalResources(getAMLocalResourcesMap(yarnConf, packagePath, configPath).asJava)
    context
  }

  private def getFs(yarnConf: YarnConfiguration) = YarnFileSystem.get(yarnConf)

  private def getAppEnv(yarnConf: YarnConfiguration): Map[String, String] = {
    val classPaths = yarnConf.getStrings(
      YarnConfiguration.YARN_APPLICATION_CLASSPATH,
      YarnConfiguration.DEFAULT_YARN_APPLICATION_CLASSPATH.mkString(File.pathSeparator))
    val allPaths = Option(classPaths).getOrElse(Array(""))

    allPaths :+ Environment.PWD.$() + File.separator + "*" + File.pathSeparator

    Map(Environment.CLASSPATH.name -> allPaths.map(_.trim).mkString(File.pathSeparator))
  }

  private def getAMLocalResourcesMap(
      yarnConf: YarnConfiguration, packagePath: String, configPath: String)
    : Map[String, LocalResource] = {
    val fs = getFs(yarnConf)

    Map(
      "pack" -> newYarnAppResource(fs, new Path(packagePath),
        LocalResourceType.ARCHIVE, LocalResourceVisibility.APPLICATION),
      "conf" -> newYarnAppResource(fs, new Path(configPath),
        LocalResourceType.FILE, LocalResourceVisibility.APPLICATION))
  }

  private def newYarnAppResource(
      fs: YarnFileSystem, path: Path,
      resourceType: LocalResourceType, vis: LocalResourceVisibility): LocalResource = {
    val qualified = fs.makeQualified(path)
    val status = fs.getFileStatus(qualified)
    val resource = Records.newRecord(classOf[LocalResource])
    resource.setType(resourceType)
    resource.setVisibility(vis)
    resource.setResource(ConverterUtils.getYarnUrlFromPath(qualified))
    resource.setTimestamp(status.getModificationTime)
    resource.setSize(status.getLen)
    resource
  }

  private def getToken(yc: YarnConfiguration, packagePath: String, configPath: String)
    : ByteBuffer = {
    val credentials = UserGroupInformation.getCurrentUser.getCredentials
    val dob = new DataOutputBuffer
    val dirs = Array(new Path(packagePath), new Path(configPath))
    TokenCache.obtainTokensForNamenodes(credentials, dirs, yc)
    credentials.writeTokenStorageToStream(dob)
    ByteBuffer.wrap(dob.getData)
  }
} 
Example 112
Source File: UDPSender.scala    From censorinus   with MIT License 5 votes vote down vote up
package github.gphat.censorinus

import java.net.{InetSocketAddress,SocketException}
import java.nio.ByteBuffer
import java.nio.channels.UnresolvedAddressException
import java.nio.channels.DatagramChannel

class UDPSender(
  hostname: String = "localhost",
  port: Int = MetricSender.DEFAULT_STATSD_PORT,
  allowExceptions: Boolean = false
) extends MetricSender {

  lazy val clientSocket = DatagramChannel.open.connect(new InetSocketAddress(hostname, port))

  def send(message: ByteBuffer): Unit = {
    try {
      val _ = clientSocket.write(message)
    } catch {
      case se @ (_ : SocketException | _ : UnresolvedAddressException) => {
        // Check if we're allowing exceptions and rethrow if so. We didn't use
        // a guard on the case because then we'd need a second case to catch
        // the !allowExceptions case!
        if(allowExceptions) {
          throw se
        }
      }
    }
  }

  def shutdown: Unit = clientSocket.close
} 
Example 113
Source File: UDPSenderSpec.scala    From censorinus   with MIT License 5 votes vote down vote up
package github.gphat.censorinus

import org.scalatest._
import org.scalatest.flatspec.AnyFlatSpec
import org.scalatest.matchers.should.Matchers

import java.nio.ByteBuffer
import java.nio.channels.UnresolvedAddressException

class UDPSenderSpec extends AnyFlatSpec with Matchers {

  "UDPSender" should "emit errors" in {
    // Guessing this port won't be used? :)
    val u = new UDPSender(hostname = "127.0.0.1789", port = 8126, allowExceptions = true)
    an [UnresolvedAddressException] should be thrownBy u.send(ByteBuffer.wrap("abc".getBytes("utf-8")))
  }

  it should "swallow errors" in {
    // Guessing this port won't be used? :)
    val u = new UDPSender(hostname = "127.0.0.1789", port = 8126, allowExceptions = false)
    u.send(ByteBuffer.wrap("abc".getBytes("utf-8")))
  }
} 
Example 114
Source File: NIOLogFileWriter.scala    From scribe   with MIT License 5 votes vote down vote up
package scribe.writer.file

import java.nio.ByteBuffer
import java.nio.channels.FileChannel
import java.nio.file.{OpenOption, StandardOpenOption}

import scala.annotation.tailrec

class NIOLogFileWriter(lf: LogFile) extends LogFileWriter {
  private lazy val options: List[OpenOption] = if (lf.append) {
    List(StandardOpenOption.WRITE, StandardOpenOption.APPEND, StandardOpenOption.CREATE)
  } else {
    List(StandardOpenOption.WRITE, StandardOpenOption.TRUNCATE_EXISTING, StandardOpenOption.CREATE)
  }
  private lazy val channel: FileChannel = FileChannel.open(lf.path, options: _*)

  override def write(output: String): Unit = {
    val bytes = output.getBytes(lf.charset)
    val buffer = ByteBuffer.wrap(bytes)
    writeBuffer(buffer)
    buffer.clear()
  }

  @tailrec
  private def writeBuffer(buffer: ByteBuffer): Unit = if (buffer.hasRemaining) {
    channel.write(buffer)
    writeBuffer(buffer)
  }

  override def flush(): Unit = channel.force(false)

  override def dispose(): Unit = if (channel.isOpen) {
    channel.close()
  }
} 
Example 115
Source File: MesosTaskLaunchData.scala    From sparkoscope   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.scheduler.cluster.mesos

import java.nio.ByteBuffer

import org.apache.mesos.protobuf.ByteString

import org.apache.spark.internal.Logging


private[spark] case class MesosTaskLaunchData(
  serializedTask: ByteBuffer,
  attemptNumber: Int) extends Logging {

  def toByteString: ByteString = {
    val dataBuffer = ByteBuffer.allocate(4 + serializedTask.limit)
    dataBuffer.putInt(attemptNumber)
    dataBuffer.put(serializedTask)
    dataBuffer.rewind
    logDebug(s"ByteBuffer size: [${dataBuffer.remaining}]")
    ByteString.copyFrom(dataBuffer)
  }
}

private[spark] object MesosTaskLaunchData extends Logging {
  def fromByteString(byteString: ByteString): MesosTaskLaunchData = {
    val byteBuffer = byteString.asReadOnlyByteBuffer()
    logDebug(s"ByteBuffer size: [${byteBuffer.remaining}]")
    val attemptNumber = byteBuffer.getInt // updates the position by 4 bytes
    val serializedTask = byteBuffer.slice() // subsequence starting at the current position
    MesosTaskLaunchData(serializedTask, attemptNumber)
  }
} 
Example 116
Source File: MesosTaskLaunchDataSuite.scala    From sparkoscope   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.scheduler.cluster.mesos

import java.nio.ByteBuffer

import org.apache.spark.SparkFunSuite

class MesosTaskLaunchDataSuite extends SparkFunSuite {
  test("serialize and deserialize data must be same") {
    val serializedTask = ByteBuffer.allocate(40)
    (Range(100, 110).map(serializedTask.putInt(_)))
    serializedTask.rewind
    val attemptNumber = 100
    val byteString = MesosTaskLaunchData(serializedTask, attemptNumber).toByteString
    serializedTask.rewind
    val mesosTaskLaunchData = MesosTaskLaunchData.fromByteString(byteString)
    assert(mesosTaskLaunchData.attemptNumber == attemptNumber)
    assert(mesosTaskLaunchData.serializedTask.equals(serializedTask))
  }
} 
Example 117
Source File: KPLBasedKinesisTestUtils.scala    From sparkoscope   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.streaming.kinesis

import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets

import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer

import com.amazonaws.services.kinesis.producer.{KinesisProducer => KPLProducer, KinesisProducerConfiguration, UserRecordResult}
import com.google.common.util.concurrent.{FutureCallback, Futures}

private[kinesis] class KPLBasedKinesisTestUtils extends KinesisTestUtils {
  override protected def getProducer(aggregate: Boolean): KinesisDataGenerator = {
    if (!aggregate) {
      new SimpleDataGenerator(kinesisClient)
    } else {
      new KPLDataGenerator(regionName)
    }
  }
}


private[kinesis] class KPLDataGenerator(regionName: String) extends KinesisDataGenerator {

  private lazy val producer: KPLProducer = {
    val conf = new KinesisProducerConfiguration()
      .setRecordMaxBufferedTime(1000)
      .setMaxConnections(1)
      .setRegion(regionName)
      .setMetricsLevel("none")

    new KPLProducer(conf)
  }

  override def sendData(streamName: String, data: Seq[Int]): Map[String, Seq[(Int, String)]] = {
    val shardIdToSeqNumbers = new mutable.HashMap[String, ArrayBuffer[(Int, String)]]()
    data.foreach { num =>
      val str = num.toString
      val data = ByteBuffer.wrap(str.getBytes(StandardCharsets.UTF_8))
      val future = producer.addUserRecord(streamName, str, data)
      val kinesisCallBack = new FutureCallback[UserRecordResult]() {
        override def onFailure(t: Throwable): Unit = {} // do nothing

        override def onSuccess(result: UserRecordResult): Unit = {
          val shardId = result.getShardId
          val seqNumber = result.getSequenceNumber()
          val sentSeqNumbers = shardIdToSeqNumbers.getOrElseUpdate(shardId,
            new ArrayBuffer[(Int, String)]())
          sentSeqNumbers += ((num, seqNumber))
        }
      }
      Futures.addCallback(future, kinesisCallBack)
    }
    producer.flushSync()
    shardIdToSeqNumbers.toMap
  }
} 
Example 118
Source File: FlumeTestUtils.scala    From sparkoscope   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.streaming.flume

import java.net.{InetSocketAddress, ServerSocket}
import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets
import java.util.{List => JList}
import java.util.Collections

import scala.collection.JavaConverters._

import org.apache.avro.ipc.NettyTransceiver
import org.apache.avro.ipc.specific.SpecificRequestor
import org.apache.commons.lang3.RandomUtils
import org.apache.flume.source.avro
import org.apache.flume.source.avro.{AvroFlumeEvent, AvroSourceProtocol}
import org.jboss.netty.channel.ChannelPipeline
import org.jboss.netty.channel.socket.SocketChannel
import org.jboss.netty.channel.socket.nio.NioClientSocketChannelFactory
import org.jboss.netty.handler.codec.compression.{ZlibDecoder, ZlibEncoder}

import org.apache.spark.util.Utils
import org.apache.spark.SparkConf


  private class CompressionChannelFactory(compressionLevel: Int)
    extends NioClientSocketChannelFactory {

    override def newChannel(pipeline: ChannelPipeline): SocketChannel = {
      val encoder = new ZlibEncoder(compressionLevel)
      pipeline.addFirst("deflater", encoder)
      pipeline.addFirst("inflater", new ZlibDecoder())
      super.newChannel(pipeline)
    }
  }

} 
Example 119
Source File: NullableColumnAccessor.scala    From sparkoscope   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.sql.execution.columnar

import java.nio.{ByteBuffer, ByteOrder}

import org.apache.spark.sql.catalyst.InternalRow

private[columnar] trait NullableColumnAccessor extends ColumnAccessor {
  private var nullsBuffer: ByteBuffer = _
  private var nullCount: Int = _
  private var seenNulls: Int = 0

  private var nextNullIndex: Int = _
  private var pos: Int = 0

  abstract override protected def initialize(): Unit = {
    nullsBuffer = underlyingBuffer.duplicate().order(ByteOrder.nativeOrder())
    nullCount = ByteBufferHelper.getInt(nullsBuffer)
    nextNullIndex = if (nullCount > 0) ByteBufferHelper.getInt(nullsBuffer) else -1
    pos = 0

    underlyingBuffer.position(underlyingBuffer.position + 4 + nullCount * 4)
    super.initialize()
  }

  abstract override def extractTo(row: InternalRow, ordinal: Int): Unit = {
    if (pos == nextNullIndex) {
      seenNulls += 1

      if (seenNulls < nullCount) {
        nextNullIndex = ByteBufferHelper.getInt(nullsBuffer)
      }

      row.setNullAt(ordinal)
    } else {
      super.extractTo(row, ordinal)
    }

    pos += 1
  }

  abstract override def hasNext: Boolean = seenNulls < nullCount || super.hasNext
} 
Example 120
Source File: Serializer.scala    From monix-kafka   with Apache License 2.0 5 votes vote down vote up
package monix.kafka

import java.nio.ByteBuffer
import org.apache.kafka.common.serialization._
import org.apache.kafka.common.serialization.{Serializer => KafkaSerializer}
import org.apache.kafka.common.utils.Bytes


  type Constructor[A] = (Serializer[A]) => KafkaSerializer[A]

  implicit val forStrings: Serializer[String] =
    Serializer[String](
      className = "org.apache.kafka.common.serialization.StringSerializer",
      classType = classOf[StringSerializer]
    )

  implicit val forByteArray: Serializer[Array[Byte]] =
    Serializer[Array[Byte]](
      className = "org.apache.kafka.common.serialization.ByteArraySerializer",
      classType = classOf[ByteArraySerializer]
    )

  implicit val forByteBuffer: Serializer[ByteBuffer] =
    Serializer[ByteBuffer](
      className = "org.apache.kafka.common.serialization.ByteBufferSerializer",
      classType = classOf[ByteBufferSerializer]
    )

  implicit val forBytes: Serializer[Bytes] =
    Serializer[Bytes](
      className = "org.apache.kafka.common.serialization.BytesSerializer",
      classType = classOf[BytesSerializer]
    )

  implicit val forJavaDouble: Serializer[java.lang.Double] =
    Serializer[java.lang.Double](
      className = "org.apache.kafka.common.serialization.DoubleSerializer",
      classType = classOf[DoubleSerializer]
    )

  implicit val forJavaInteger: Serializer[java.lang.Integer] =
    Serializer[java.lang.Integer](
      className = "org.apache.kafka.common.serialization.IntegerSerializer",
      classType = classOf[IntegerSerializer]
    )

  implicit val forJavaLong: Serializer[java.lang.Long] =
    Serializer[java.lang.Long](
      className = "org.apache.kafka.common.serialization.LongSerializer",
      classType = classOf[LongSerializer]
    )
} 
Example 121
Source File: ByteBufferChopper.scala    From RosHTTP   with MIT License 5 votes vote down vote up
package fr.hmil.roshttp

import java.nio.ByteBuffer


object ByteBufferChopper {

  def chop(buffer: ByteBuffer, maxChunkSize: Int): Seq[ByteBuffer] = {
    val nb_buffers = (buffer.limit() + maxChunkSize - 1) / maxChunkSize
    val buffers = new Array[ByteBuffer](nb_buffers)
    var i = 0
    while (i < nb_buffers) {
      val length = Math.min(maxChunkSize, buffer.remaining)
      buffers(i) = buffer.slice()
      buffers(i).limit(length)
      buffer.position(buffer.position() + length)
      i = i + 1
    }
    buffers
  }

} 
Example 122
Source File: Converters.scala    From RosHTTP   with MIT License 5 votes vote down vote up
package fr.hmil.roshttp

import java.nio.ByteBuffer

import fr.hmil.roshttp.node.buffer.Buffer

import scala.scalajs.js
import scala.scalajs.js.typedarray.{ArrayBuffer, Int8Array, TypedArrayBuffer, Uint8Array}
import scala.scalajs.js.JSConverters._
import js.typedarray.TypedArrayBufferOps._

private object Converters {
  def byteArrayToUint8Array(arr: Array[Byte]): Uint8Array = {
    js.Dynamic.newInstance(js.Dynamic.global.Uint8Array)(arr.toJSArray).asInstanceOf[Uint8Array]
  }

  def byteBufferToNodeBuffer(buffer: ByteBuffer): Buffer = {
    if (buffer.isDirect) {
      js.Dynamic.newInstance(js.Dynamic.global.Buffer)(buffer.arrayBuffer).asInstanceOf[Buffer]
    } else if (buffer.hasArray) {
      js.Dynamic.newInstance(js.Dynamic.global.Buffer)(byteArrayToUint8Array(buffer.array())).asInstanceOf[Buffer]
    } else {
      val arr = new Int8Array(buffer.limit())
      var i = 0
      while (i < arr.length) {
        arr(i) = buffer.get(i)
        i += 1
      }
      js.Dynamic.newInstance(js.Dynamic.global.Buffer)(arr).asInstanceOf[Buffer]
    }
  }

  def nodeBufferToByteBuffer(buffer: Buffer): ByteBuffer = {
    TypedArrayBuffer.wrap(buffer.asInstanceOf[ArrayBuffer])
  }

  def arrayBufferToByteBuffer(buffer: ArrayBuffer): ByteBuffer = {
    TypedArrayBuffer.wrap(buffer)
  }
} 
Example 123
Source File: Utils.scala    From RosHTTP   with MIT License 5 votes vote down vote up
package fr.hmil.roshttp.util

import java.nio.ByteBuffer

import fr.hmil.roshttp.CrossPlatformUtils

object Utils {

  
  def encodeQueryString(queryString: String): String = {
    queryString
      .split("&")
      .map(_
        .split("=")
        .map(encodeURIComponent)
        .mkString("="))
      .mkString("&")
  }

  def encodeURIComponent(input: String): String = CrossPlatformUtils.encodeURIComponent(input)

  def getStringFromBuffer(byteBuffer: ByteBuffer, charset: String): String = {
    if (byteBuffer.hasArray) {
      new String(byteBuffer.array(), 0, byteBuffer.limit(), charset)
    } else {
      val tmp = new Array[Byte](byteBuffer.limit)
      byteBuffer.get(tmp)
      new String(tmp, charset)
    }
  }

  private val oneByteCharset = "utf-8"
} 
Example 124
Source File: StreamHttpResponse.scala    From RosHTTP   with MIT License 5 votes vote down vote up
package fr.hmil.roshttp.response

import java.nio.ByteBuffer

import fr.hmil.roshttp.BackendConfig
import fr.hmil.roshttp.util.HeaderMap
import monix.execution.Scheduler
import monix.reactive.Observable

import scala.concurrent.Future


class StreamHttpResponse(
    val statusCode: Int,
    val headers: HeaderMap[String],
    val body: Observable[ByteBuffer])
extends HttpResponse

object StreamHttpResponse extends HttpResponseFactory[StreamHttpResponse] {
  override def apply(
      header: HttpResponseHeader,
      bodyStream: Observable[ByteBuffer],
      config: BackendConfig)
      (implicit scheduler: Scheduler): Future[StreamHttpResponse] =
    Future.successful(new StreamHttpResponse(header.statusCode, header.headers, bodyStream))
} 
Example 125
Source File: SimpleHttpResponse.scala    From RosHTTP   with MIT License 5 votes vote down vote up
package fr.hmil.roshttp.response

import java.nio.ByteBuffer

import fr.hmil.roshttp.BackendConfig
import fr.hmil.roshttp.exceptions.ResponseException
import fr.hmil.roshttp.util.{HeaderMap, Utils}
import monix.execution.Scheduler
import monix.reactive.Observable

import scala.collection.mutable
import scala.concurrent.{Future, Promise}
import scala.util.{Failure, Success}


class SimpleHttpResponse(
    val statusCode: Int,
    val headers: HeaderMap[String],
    val body: String)
  extends HttpResponse

object SimpleHttpResponse extends HttpResponseFactory[SimpleHttpResponse] {
  override def apply(
      header: HttpResponseHeader,
      bodyStream: Observable[ByteBuffer],
      config: BackendConfig)
      (implicit scheduler: Scheduler): Future[SimpleHttpResponse] = {

    val charset = Utils.charsetFromContentType(header.headers.getOrElse("content-type", null))
    val buffers = mutable.Queue[ByteBuffer]()
    val promise = Promise[SimpleHttpResponse]()

    val streamCollector = bodyStream.
      foreach(elem => buffers.enqueue(elem)).
      map({_ =>
        val body = recomposeBody(buffers, config.maxChunkSize, charset)
        new SimpleHttpResponse(header.statusCode, header.headers, body)
      })

    streamCollector.onComplete({
      case res:Success[SimpleHttpResponse] =>
        promise.trySuccess(res.value)
      case e:Failure[_] =>
        promise.tryFailure(new ResponseException(e.exception, header))
    })

    promise.future
  }

  private def recomposeBody(seq: mutable.Queue[ByteBuffer], maxChunkSize: Int, charset: String): String = {
    // Allocate maximum expected body length
    val buffer = ByteBuffer.allocate(seq.length * maxChunkSize)
    val totalBytes = seq.foldLeft(0)({ (count, chunk) =>
      buffer.put(chunk)
      count + chunk.limit()
    })
    buffer.limit(totalBytes)
    Utils.getStringFromBuffer(buffer, charset)
  }
} 
Example 126
Source File: ByteBufferBody.scala    From RosHTTP   with MIT License 5 votes vote down vote up
package fr.hmil.roshttp.body

import java.nio.ByteBuffer


class ByteBufferBody private(
    data: ByteBuffer,
    override val contentType: String
  ) extends BulkBodyPart {
  override def contentData: ByteBuffer = data
}

object ByteBufferBody {
  def apply(data: ByteBuffer, contentType: String = "application/octet-stream"): ByteBufferBody =
    new ByteBufferBody(data, contentType)
} 
Example 127
Source File: MultiPartBody.scala    From RosHTTP   with MIT License 5 votes vote down vote up
package fr.hmil.roshttp.body

import java.nio.ByteBuffer

import monix.execution.Scheduler
import monix.reactive.Observable

import scala.util.Random


class MultiPartBody(parts: Map[String, BodyPart], subtype: String = "form-data")(implicit scheduler: Scheduler)
  extends BodyPart {

  val boundary = "----" + Random.alphanumeric.take(24).mkString.toLowerCase

  override def contentType: String = s"multipart/$subtype; boundary=$boundary"

  override def content: Observable[ByteBuffer] = {
    parts.
      // Prepend multipart encapsulation boundary and body part headers to
      // each body part.
      map({ case (name, part) =>
        ByteBuffer.wrap(
          ("\r\n--" + boundary + "\r\n" +
            "Content-Disposition: form-data; name=\"" + name + "\"\r\n" +
            s"Content-Type: ${part.contentType}\r\n" +
            "\r\n").getBytes("utf-8")
        ) +: part.content
      }).
      // Join body parts
      reduceLeft((acc, elem) => acc ++ elem).
      // Append the closing boundary
      :+(ByteBuffer.wrap(s"\r\n--$boundary--\r\n".getBytes("utf-8")))
  }
}

object MultiPartBody {
  def apply(parts: (String, BodyPart)*)(implicit scheduler: Scheduler): MultiPartBody =
    new MultiPartBody(Map(parts: _*))
} 
Example 128
Source File: JSONBody.scala    From RosHTTP   with MIT License 5 votes vote down vote up
package fr.hmil.roshttp.body

import java.nio.ByteBuffer

import fr.hmil.roshttp.body.JSONBody.JSONValue


    while (i != end) {
      val start = i
      var c: Int = str.charAt(i)
      // Find all consecutive ASCII printable characters from `start`
      while (i != end && c >= 32 && c <= 126 && c != 34 && c != 92) {
        i += 1
        if (i != end)
          c = str.charAt(i)
      }
      // Print ASCII printable characters from `start`
      if (start != i)
        out.append(str, start, i)

      // Print next non ASCII printable character
      if (i != end) {
        def escapeJSEncoded(c: Int): Unit = {
          if (6 < c && c < 14) {
            val i = 2 * (c - 7)
            out.append(EscapeJSChars, i, i + 2)
          } else if (c == 34) {
            out.append(EscapeJSChars, 14, 16)
          } else if (c == 92) {
            out.append(EscapeJSChars, 16, 18)
          } else {
            out.append(f"\\u$c%04x")
          }
        }
        escapeJSEncoded(c)
        i += 1
      }
    }
  }
} 
Example 129
Source File: Implicits.scala    From RosHTTP   with MIT License 5 votes vote down vote up
package fr.hmil.roshttp.body

import java.io.InputStream
import java.nio.ByteBuffer

import fr.hmil.roshttp.body.JSONBody._
import monix.reactive.Observable
import monix.eval.Task


object Implicits {
  implicit def stringToJSONString(value: String): JSONString = new JSONString(value)
  implicit def intToJSONNumber(value: Int): JSONNumber = new JSONNumber(value)
  implicit def floatToJSONNumber(value: Float): JSONNumber = new JSONNumber(value)
  implicit def doubleToJSONNumber(value: Double): JSONNumber = new JSONNumber(value)
  implicit def booleanToJSONBoolean(value: Boolean):JSONBoolean = new JSONBoolean(value)
  implicit def JSONObjectToJSONBody(obj: JSONObject): JSONBody = JSONBody(obj)
  implicit def JSONArrayToJSONBody(arr: JSONArray): JSONBody = JSONBody(arr)

  implicit def byteBufferToByteBufferBody(buffer: ByteBuffer): BodyPart = ByteBufferBody(buffer)
  implicit def observableToStreamBody(is: InputStream): BodyPart =
    StreamBody(Observable.fromInputStream(Task(is)).map(ByteBuffer.wrap))
} 
Example 130
Source File: Header.scala    From spark-bam   with Apache License 2.0 5 votes vote down vote up
package org.hammerlab.bgzf.block

import java.io.{ IOException, InputStream }
import java.nio.ByteBuffer

import org.hammerlab.channel.ByteChannel


case class Header(size: Int, compressedSize: Int)

object Header {

  // 18 bytes is enough to learn what we need to know: sizes of header and compressed block
  val EXPECTED_HEADER_SIZE = 18

  def apply(ch: ByteChannel)(implicit buf: ByteBuffer): Header = {
    buf.limit(EXPECTED_HEADER_SIZE)
    ch.readFully(buf)

    implicit val arr = buf.array
    val header = make
    buf.clear()
    ch.skip(header.size - EXPECTED_HEADER_SIZE)
    buf.position(header.size)

    header
  }

  def apply(is: InputStream)(implicit buffer: Array[Byte]): Header = {

    val headerBytesRead = is.read(buffer, 0, EXPECTED_HEADER_SIZE)
    if (headerBytesRead != EXPECTED_HEADER_SIZE)
      throw new IOException(
        s"Expected $EXPECTED_HEADER_SIZE header bytes, got $headerBytesRead"
      )

    val header = make
    is.skip(header.size - EXPECTED_HEADER_SIZE)

    header
  }

  def make(implicit bytes: Array[Byte]): Header = {

    def check(idx: Int, expected: Byte): Unit = {
      val actual = bytes(idx)
      if (actual != expected)
        throw HeaderParseException(
          idx,
          actual,
          expected
        )
    }

    // GZip magic bytes
    check(0,  31)
    check(1, 139.toByte)
    check(2,   8)
    check(3,   4)

    val xlen = getShort(10)

    // We expect 6 bytes of `xlen`; anything more is considered "extra" and added to the expected 18-byte header size
    val extraHeaderBytes = xlen - 6
    val actualHeaderSize = EXPECTED_HEADER_SIZE + extraHeaderBytes

    // BAM-specific GZip-flags
    check(12, 66)
    check(13, 67)
    check(14,  2)

    val compressedSize = getShort(16) + 1

    Header(
      actualHeaderSize,
      compressedSize
    )
  }

  def getShort(idx: Int)(implicit buffer: Array[Byte]): Int =
    (buffer(idx) & 0xff) |
      ((buffer(idx + 1) & 0xff) << 8)
} 
Example 131
Source File: Deserializer.scala    From monix-kafka   with Apache License 2.0 5 votes vote down vote up
package monix.kafka

import java.nio.ByteBuffer
import org.apache.kafka.common.serialization._
import org.apache.kafka.common.serialization.{Deserializer => KafkaDeserializer}
import org.apache.kafka.common.utils.Bytes


  type Constructor[A] = (Deserializer[A]) => KafkaDeserializer[A]

  implicit val forStrings: Deserializer[String] =
    Deserializer[String](
      className = "org.apache.kafka.common.serialization.StringDeserializer",
      classType = classOf[StringDeserializer]
    )

  implicit val forByteArray: Deserializer[Array[Byte]] =
    Deserializer[Array[Byte]](
      className = "org.apache.kafka.common.serialization.ByteArrayDeserializer",
      classType = classOf[ByteArrayDeserializer]
    )

  implicit val forByteBuffer: Deserializer[ByteBuffer] =
    Deserializer[ByteBuffer](
      className = "org.apache.kafka.common.serialization.ByteBufferDeserializer",
      classType = classOf[ByteBufferDeserializer]
    )

  implicit val forBytes: Deserializer[Bytes] =
    Deserializer[Bytes](
      className = "org.apache.kafka.common.serialization.BytesDeserializer",
      classType = classOf[BytesDeserializer]
    )

  implicit val forJavaDouble: Deserializer[java.lang.Double] =
    Deserializer[java.lang.Double](
      className = "org.apache.kafka.common.serialization.DoubleDeserializer",
      classType = classOf[DoubleDeserializer]
    )

  implicit val forJavaInteger: Deserializer[java.lang.Integer] =
    Deserializer[java.lang.Integer](
      className = "org.apache.kafka.common.serialization.IntegerDeserializer",
      classType = classOf[IntegerDeserializer]
    )

  implicit val forJavaLong: Deserializer[java.lang.Long] =
    Deserializer[java.lang.Long](
      className = "org.apache.kafka.common.serialization.LongDeserializer",
      classType = classOf[LongDeserializer]
    )
} 
Example 132
Source File: PgWireProtocolSuite.scala    From spark-sql-server   with Apache License 2.0 5 votes vote down vote up
package org.apache.spark.sql.server.service.postgresql.protocol.v3

import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets
import java.sql.SQLException

import org.apache.spark.SparkFunSuite
import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.StructType
import org.apache.spark.unsafe.types.UTF8String

class PgWireProtocolSuite extends SparkFunSuite {

  val conf = new SQLConf()

  test("DataRow") {
    val v3Protocol = new PgWireProtocol(65536)
    val row = new GenericInternalRow(2)
    row.update(0, 8)
    row.update(1, UTF8String.fromString("abcdefghij"))
    val schema = StructType.fromDDL("a INT, b STRING")
    val rowConverters = PgRowConverters(conf, schema, Seq(true, false))
    val data = v3Protocol.DataRow(row, rowConverters)
    val bytes = ByteBuffer.wrap(data)
    assert(bytes.get() === 'D'.toByte)
    assert(bytes.getInt === 28)
    assert(bytes.getShort === 2)
    assert(bytes.getInt === 4)
    assert(bytes.getInt === 8)
    assert(bytes.getInt === 10)
    assert(data.slice(19, 30) === "abcdefghij".getBytes(StandardCharsets.UTF_8))
  }

  test("Fails when message buffer overflowed") {
    val v3Protocol = new PgWireProtocol(4)
    val row = new GenericInternalRow(1)
    row.update(0, UTF8String.fromString("abcdefghijk"))
    val schema = StructType.fromDDL("a STRING")
    val rowConverters = PgRowConverters(conf, schema, Seq(false))
    val errMsg = intercept[SQLException] {
      v3Protocol.DataRow(row, rowConverters)
    }.getMessage
    assert(errMsg.contains(
      "Cannot generate a V3 protocol message because buffer is not enough for the message. " +
        "To avoid this exception, you might set higher value at " +
        "'spark.sql.server.messageBufferSizeInBytes'")
    )
  }
} 
Example 133
Source File: Deserializer.scala    From monix-kafka   with Apache License 2.0 5 votes vote down vote up
package monix.kafka

import java.nio.ByteBuffer
import org.apache.kafka.common.serialization._
import org.apache.kafka.common.serialization.{Deserializer => KafkaDeserializer}
import org.apache.kafka.common.utils.Bytes


  type Constructor[A] = (Deserializer[A]) => KafkaDeserializer[A]

  implicit val forStrings: Deserializer[String] =
    Deserializer[String](
      className = "org.apache.kafka.common.serialization.StringDeserializer",
      classType = classOf[StringDeserializer]
    )

  implicit val forByteArray: Deserializer[Array[Byte]] =
    Deserializer[Array[Byte]](
      className = "org.apache.kafka.common.serialization.ByteArrayDeserializer",
      classType = classOf[ByteArrayDeserializer]
    )

  implicit val forByteBuffer: Deserializer[ByteBuffer] =
    Deserializer[ByteBuffer](
      className = "org.apache.kafka.common.serialization.ByteBufferDeserializer",
      classType = classOf[ByteBufferDeserializer]
    )

  implicit val forBytes: Deserializer[Bytes] =
    Deserializer[Bytes](
      className = "org.apache.kafka.common.serialization.BytesDeserializer",
      classType = classOf[BytesDeserializer]
    )

  implicit val forJavaDouble: Deserializer[java.lang.Double] =
    Deserializer[java.lang.Double](
      className = "org.apache.kafka.common.serialization.DoubleDeserializer",
      classType = classOf[DoubleDeserializer]
    )

  implicit val forJavaInteger: Deserializer[java.lang.Integer] =
    Deserializer[java.lang.Integer](
      className = "org.apache.kafka.common.serialization.IntegerDeserializer",
      classType = classOf[IntegerDeserializer]
    )

  implicit val forJavaLong: Deserializer[java.lang.Long] =
    Deserializer[java.lang.Long](
      className = "org.apache.kafka.common.serialization.LongDeserializer",
      classType = classOf[LongDeserializer]
    )
} 
Example 134
Source File: Serializer.scala    From monix-kafka   with Apache License 2.0 5 votes vote down vote up
package monix.kafka

import java.nio.ByteBuffer
import org.apache.kafka.common.serialization._
import org.apache.kafka.common.serialization.{Serializer => KafkaSerializer}
import org.apache.kafka.common.utils.Bytes


  type Constructor[A] = (Serializer[A]) => KafkaSerializer[A]

  implicit val forStrings: Serializer[String] =
    Serializer[String](
      className = "org.apache.kafka.common.serialization.StringSerializer",
      classType = classOf[StringSerializer]
    )

  implicit val forByteArray: Serializer[Array[Byte]] =
    Serializer[Array[Byte]](
      className = "org.apache.kafka.common.serialization.ByteArraySerializer",
      classType = classOf[ByteArraySerializer]
    )

  implicit val forByteBuffer: Serializer[ByteBuffer] =
    Serializer[ByteBuffer](
      className = "org.apache.kafka.common.serialization.ByteBufferSerializer",
      classType = classOf[ByteBufferSerializer]
    )

  implicit val forBytes: Serializer[Bytes] =
    Serializer[Bytes](
      className = "org.apache.kafka.common.serialization.BytesSerializer",
      classType = classOf[BytesSerializer]
    )

  implicit val forJavaDouble: Serializer[java.lang.Double] =
    Serializer[java.lang.Double](
      className = "org.apache.kafka.common.serialization.DoubleSerializer",
      classType = classOf[DoubleSerializer]
    )

  implicit val forJavaInteger: Serializer[java.lang.Integer] =
    Serializer[java.lang.Integer](
      className = "org.apache.kafka.common.serialization.IntegerSerializer",
      classType = classOf[IntegerSerializer]
    )

  implicit val forJavaLong: Serializer[java.lang.Long] =
    Serializer[java.lang.Long](
      className = "org.apache.kafka.common.serialization.LongSerializer",
      classType = classOf[LongSerializer]
    )
} 
Example 135
Source File: Deserializer.scala    From monix-kafka   with Apache License 2.0 5 votes vote down vote up
package monix.kafka

import java.nio.ByteBuffer
import org.apache.kafka.common.serialization._
import org.apache.kafka.common.serialization.{Deserializer => KafkaDeserializer}
import org.apache.kafka.common.utils.Bytes


  type Constructor[A] = (Deserializer[A]) => KafkaDeserializer[A]

  implicit val forStrings: Deserializer[String] =
    Deserializer[String](
      className = "org.apache.kafka.common.serialization.StringDeserializer",
      classType = classOf[StringDeserializer]
    )

  implicit val forByteArray: Deserializer[Array[Byte]] =
    Deserializer[Array[Byte]](
      className = "org.apache.kafka.common.serialization.ByteArrayDeserializer",
      classType = classOf[ByteArrayDeserializer]
    )

  implicit val forByteBuffer: Deserializer[ByteBuffer] =
    Deserializer[ByteBuffer](
      className = "org.apache.kafka.common.serialization.ByteBufferDeserializer",
      classType = classOf[ByteBufferDeserializer]
    )

  implicit val forBytes: Deserializer[Bytes] =
    Deserializer[Bytes](
      className = "org.apache.kafka.common.serialization.BytesDeserializer",
      classType = classOf[BytesDeserializer]
    )

  implicit val forJavaDouble: Deserializer[java.lang.Double] =
    Deserializer[java.lang.Double](
      className = "org.apache.kafka.common.serialization.DoubleDeserializer",
      classType = classOf[DoubleDeserializer]
    )

  implicit val forJavaInteger: Deserializer[java.lang.Integer] =
    Deserializer[java.lang.Integer](
      className = "org.apache.kafka.common.serialization.IntegerDeserializer",
      classType = classOf[IntegerDeserializer]
    )

  implicit val forJavaLong: Deserializer[java.lang.Long] =
    Deserializer[java.lang.Long](
      className = "org.apache.kafka.common.serialization.LongDeserializer",
      classType = classOf[LongDeserializer]
    )
} 
Example 136
Source File: Serializer.scala    From monix-kafka   with Apache License 2.0 5 votes vote down vote up
package monix.kafka

import java.nio.ByteBuffer
import org.apache.kafka.common.serialization._
import org.apache.kafka.common.serialization.{Serializer => KafkaSerializer}
import org.apache.kafka.common.utils.Bytes


  type Constructor[A] = (Serializer[A]) => KafkaSerializer[A]

  implicit val forStrings: Serializer[String] =
    Serializer[String](
      className = "org.apache.kafka.common.serialization.StringSerializer",
      classType = classOf[StringSerializer]
    )

  implicit val forByteArray: Serializer[Array[Byte]] =
    Serializer[Array[Byte]](
      className = "org.apache.kafka.common.serialization.ByteArraySerializer",
      classType = classOf[ByteArraySerializer]
    )

  implicit val forByteBuffer: Serializer[ByteBuffer] =
    Serializer[ByteBuffer](
      className = "org.apache.kafka.common.serialization.ByteBufferSerializer",
      classType = classOf[ByteBufferSerializer]
    )

  implicit val forBytes: Serializer[Bytes] =
    Serializer[Bytes](
      className = "org.apache.kafka.common.serialization.BytesSerializer",
      classType = classOf[BytesSerializer]
    )

  implicit val forJavaDouble: Serializer[java.lang.Double] =
    Serializer[java.lang.Double](
      className = "org.apache.kafka.common.serialization.DoubleSerializer",
      classType = classOf[DoubleSerializer]
    )

  implicit val forJavaInteger: Serializer[java.lang.Integer] =
    Serializer[java.lang.Integer](
      className = "org.apache.kafka.common.serialization.IntegerSerializer",
      classType = classOf[IntegerSerializer]
    )

  implicit val forJavaLong: Serializer[java.lang.Long] =
    Serializer[java.lang.Long](
      className = "org.apache.kafka.common.serialization.LongSerializer",
      classType = classOf[LongSerializer]
    )
} 
Example 137
Source File: isabelle_charset.scala    From libisabelle   with Apache License 2.0 5 votes vote down vote up
package isabelle


import java.nio.Buffer
import java.nio.{ByteBuffer, CharBuffer}
import java.nio.charset.{Charset, CharsetDecoder, CharsetEncoder, CoderResult}
import java.nio.charset.spi.CharsetProvider


object Isabelle_Charset
{
  val name: String = "UTF-8-Isabelle-test"  // FIXME
  lazy val charset: Charset = new Isabelle_Charset
}


class Isabelle_Charset extends Charset(Isabelle_Charset.name, null)
{
  override def contains(cs: Charset): Boolean =
    cs.name.equalsIgnoreCase(UTF8.charset_name) || UTF8.charset.contains(cs)

  override def newDecoder(): CharsetDecoder = UTF8.charset.newDecoder

  override def newEncoder(): CharsetEncoder = UTF8.charset.newEncoder
}


class Isabelle_Charset_Provider extends CharsetProvider
{
  override def charsetForName(name: String): Charset =
  {
    // FIXME inactive
    // if (name.equalsIgnoreCase(Isabelle_Charset.name)) Isabelle_Charset.charset
    // else null
    null
  }

  override def charsets(): java.util.Iterator[Charset] =
  {
    import scala.collection.JavaConversions._
    // FIXME inactive
    // Iterator(Isabelle_Charset.charset)
    Iterator()
  }
} 
Example 138
Source File: isabelle_charset.scala    From libisabelle   with Apache License 2.0 5 votes vote down vote up
package isabelle


import java.nio.Buffer
import java.nio.{ByteBuffer, CharBuffer}
import java.nio.charset.{Charset, CharsetDecoder, CharsetEncoder, CoderResult}
import java.nio.charset.spi.CharsetProvider


object Isabelle_Charset
{
  val name: String = "UTF-8-Isabelle-test"  // FIXME
  lazy val charset: Charset = new Isabelle_Charset
}


class Isabelle_Charset extends Charset(Isabelle_Charset.name, null)
{
  override def contains(cs: Charset): Boolean =
    cs.name.equalsIgnoreCase(UTF8.charset_name) || UTF8.charset.contains(cs)

  override def newDecoder(): CharsetDecoder = UTF8.charset.newDecoder

  override def newEncoder(): CharsetEncoder = UTF8.charset.newEncoder
}


class Isabelle_Charset_Provider extends CharsetProvider
{
  override def charsetForName(name: String): Charset =
  {
    // FIXME inactive
    // if (name.equalsIgnoreCase(Isabelle_Charset.name)) Isabelle_Charset.charset
    // else null
    null
  }

  override def charsets(): java.util.Iterator[Charset] =
  {
    import scala.collection.JavaConversions._
    // FIXME inactive
    // Iterator(Isabelle_Charset.charset)
    Iterator()
  }
} 
Example 139
Source File: isabelle_charset.scala    From libisabelle   with Apache License 2.0 5 votes vote down vote up
package isabelle


import java.nio.Buffer
import java.nio.{ByteBuffer, CharBuffer}
import java.nio.charset.{Charset, CharsetDecoder, CharsetEncoder, CoderResult}
import java.nio.charset.spi.CharsetProvider


object Isabelle_Charset
{
  val name: String = "UTF-8-Isabelle-test"  // FIXME
  lazy val charset: Charset = new Isabelle_Charset
}


class Isabelle_Charset extends Charset(Isabelle_Charset.name, null)
{
  override def contains(cs: Charset): Boolean =
    cs.name.equalsIgnoreCase(UTF8.charset_name) || UTF8.charset.contains(cs)

  override def newDecoder(): CharsetDecoder = UTF8.charset.newDecoder

  override def newEncoder(): CharsetEncoder = UTF8.charset.newEncoder
}


class Isabelle_Charset_Provider extends CharsetProvider
{
  override def charsetForName(name: String): Charset =
  {
    // FIXME inactive
    // if (name.equalsIgnoreCase(Isabelle_Charset.name)) Isabelle_Charset.charset
    // else null
    null
  }

  override def charsets(): java.util.Iterator[Charset] =
  {
    import scala.collection.JavaConversions._
    // FIXME inactive
    // Iterator(Isabelle_Charset.charset)
    Iterator()
  }
} 
Example 140
Source File: AvroSerde.scala    From event-sourcing-kafka-streams   with MIT License 5 votes vote down vote up
package org.amitayh.invoices.common.serde

import java.io.ByteArrayOutputStream
import java.nio.ByteBuffer
import java.time.Instant
import java.util
import java.util.UUID

import com.sksamuel.avro4s._
import org.amitayh.invoices.common.domain._
import org.amitayh.invoices.common.serde.UuidConverters.{fromByteBuffer, toByteBuffer}
import org.apache.avro.Schema
import org.apache.avro.Schema.Field
import org.apache.kafka.common.serialization.{Deserializer, Serde, Serializer}

object AvroSerde {
  implicit val instantToSchema: ToSchema[Instant] = new ToSchema[Instant] {
    override val schema: Schema = Schema.create(Schema.Type.STRING)
  }

  implicit val instantToValue: ToValue[Instant] = new ToValue[Instant] {
    override def apply(value: Instant): String = value.toString
  }

  implicit val instantFromValue: FromValue[Instant] = new FromValue[Instant] {
    override def apply(value: Any, field: Field): Instant =
      Instant.parse(value.toString)
  }

  implicit val uuidToSchema: ToSchema[UUID] = new ToSchema[UUID] {
    override val schema: Schema = Schema.create(Schema.Type.BYTES)
  }

  implicit val uuidToValue: ToValue[UUID] = new ToValue[UUID] {
    override def apply(value: UUID): ByteBuffer = toByteBuffer(value)
  }

  implicit val uuidFromValue: FromValue[UUID] = new FromValue[UUID] {
    override def apply(value: Any, field: Field): UUID =
      fromByteBuffer(value.asInstanceOf[ByteBuffer])
  }

  val CommandSerde: Serde[Command] = serdeFor[Command]

  val CommandResultSerde: Serde[CommandResult] = serdeFor[CommandResult]

  val SnapshotSerde: Serde[InvoiceSnapshot] = serdeFor[InvoiceSnapshot]

  val EventSerde: Serde[Event] = serdeFor[Event]

  def toBytes[T: SchemaFor: ToRecord](data: T): Array[Byte] = {
    val baos = new ByteArrayOutputStream
    val output = AvroOutputStream.binary[T](baos)
    output.write(data)
    output.close()
    baos.toByteArray
  }

  def fromBytes[T: SchemaFor: FromRecord](data: Array[Byte]): T = {
    val input = AvroInputStream.binary[T](data)
    input.iterator.next()
  }

  private def serdeFor[T: SchemaFor: ToRecord: FromRecord]: Serde[T] = new Serde[T] {
    override val serializer: Serializer[T] = new Serializer[T] {
      override def serialize(topic: String, data: T): Array[Byte] = toBytes(data)
      override def configure(configs: util.Map[String, _], isKey: Boolean): Unit = ()
      override def close(): Unit = ()
    }
    override val deserializer: Deserializer[T] = new Deserializer[T] {
      override def deserialize(topic: String, data: Array[Byte]): T = fromBytes(data)
      override def configure(configs: util.Map[String, _], isKey: Boolean): Unit = ()
      override def close(): Unit = ()
    }
    override def configure(configs: util.Map[String, _], isKey: Boolean): Unit = ()
    override def close(): Unit = ()
  }
} 
Example 141
Source File: UuidConverters.scala    From event-sourcing-kafka-streams   with MIT License 5 votes vote down vote up
package org.amitayh.invoices.common.serde

import java.nio.ByteBuffer
import java.util.UUID

object UuidConverters {
  def toBytes(uuid: UUID): Array[Byte] =
    toByteBuffer(uuid).array

  def toByteBuffer(uuid: UUID): ByteBuffer = {
    val buffer = ByteBuffer.allocate(16)
    buffer.putLong(0, uuid.getMostSignificantBits)
    buffer.putLong(8, uuid.getLeastSignificantBits)
    buffer
  }

  def fromBytes(data: Array[Byte]): UUID =
    fromByteBuffer(ByteBuffer.wrap(data))

  def fromByteBuffer(buffer: ByteBuffer): UUID = {
    val mostSignificantBits = buffer.getLong(0)
    val leastSignificantBits = buffer.getLong(8)
    new UUID(mostSignificantBits, leastSignificantBits)
  }
} 
Example 142
Source File: FinatraRequestToRawBody.scala    From tapir   with Apache License 2.0 5 votes vote down vote up
package sttp.tapir.server.finatra

import java.io.ByteArrayInputStream
import java.nio.ByteBuffer
import java.nio.charset.Charset

import com.twitter.finagle.http.Request
import com.twitter.finatra.http.request.RequestUtils
import com.twitter.io.Buf
import com.twitter.util.Future
import org.apache.commons.fileupload.FileItemHeaders
import sttp.model.{Part, Header}
import sttp.tapir.{RawPart, RawBodyType}

import scala.collection.immutable.Seq
import scala.collection.JavaConverters._

class FinatraRequestToRawBody(serverOptions: FinatraServerOptions) {
  def apply[R](bodyType: RawBodyType[R], body: Buf, charset: Option[Charset], request: Request): Future[R] = {
    def asByteArray: Array[Byte] = {
      val array = new Array[Byte](body.length)
      body.write(array, 0)
      array
    }

    def asByteBuffer: ByteBuffer = {
      val buffer = ByteBuffer.allocate(body.length)
      body.write(buffer)
      buffer.flip()
      buffer
    }

    bodyType match {
      case RawBodyType.StringBody(defaultCharset) => Future.value[R](new String(asByteArray, charset.getOrElse(defaultCharset)))
      case RawBodyType.ByteArrayBody              => Future.value[R](asByteArray)
      case RawBodyType.ByteBufferBody             => Future.value[R](asByteBuffer)
      case RawBodyType.InputStreamBody            => Future.value[R](new ByteArrayInputStream(asByteArray))
      case RawBodyType.FileBody                   => serverOptions.createFile(asByteArray)
      case m: RawBodyType.MultipartBody           => multiPartRequestToRawBody(request, m)
    }
  }

  private def parseDispositionParams(headerValue: Option[String]): Map[String, String] =
    headerValue
      .map(
        _.split(";")
          .map(_.trim)
          .tail
          .map(_.split("="))
          .map(array => array(0) -> array(1))
          .toMap
      )
      .getOrElse(Map.empty)

  private def getCharset(contentType: Option[String]): Option[Charset] =
    contentType.flatMap(
      _.split(";")
        .map(_.trim)
        .tail
        .map(_.split("="))
        .map(array => array(0) -> array(1))
        .toMap
        .get("charset")
        .map(Charset.forName)
    )

  private def multiPartRequestToRawBody(request: Request, m: RawBodyType.MultipartBody): Future[Seq[RawPart]] = {
    def fileItemHeaders(headers: FileItemHeaders): Seq[Header] = {
      headers.getHeaderNames.asScala
        .flatMap { name => headers.getHeaders(name).asScala.map(name -> _) }
        .toSeq
        .filter(_._1.toLowerCase != "content-disposition")
        .map { case (k, v) => Header(k, v) }
        .toList
    }

    Future
      .collect(
        RequestUtils
          .multiParams(request)
          .flatMap {
            case (name, multiPartItem) =>
              val dispositionParams: Map[String, String] =
                parseDispositionParams(Option(multiPartItem.headers.getHeader("content-disposition")))
              val charset = getCharset(multiPartItem.contentType)

              for {
                partType <- m.partType(name)
                futureBody = apply(partType, Buf.ByteArray.Owned(multiPartItem.data), charset, request)
              } yield futureBody
                .map(body =>
                  Part(name, body, otherDispositionParams = dispositionParams - "name", headers = fileItemHeaders(multiPartItem.headers))
                    .asInstanceOf[RawPart]
                )
          }
          .toSeq
      )
      .map(_.toList)
  }
} 
Example 143
Source File: ByteBuffers.scala    From cosmos   with Apache License 2.0 5 votes vote down vote up
package com.mesosphere.universe.common

import java.nio.ByteBuffer

object ByteBuffers {

  def getBytes(bb: ByteBuffer): Array[Byte] = {
    if (bb.hasArray){
      bb.array()
    } else {
      //the index of the first byte in the backing byte array of this ByteBuffer.
      //if there isn't a backing byte array, then this throws
      val position = bb.position()

      val remaining = bb.remaining()
      if (position > remaining) {
        throw new IndexOutOfBoundsException(s"$position > $remaining")
      } else {
        val dest = new Array[Byte](remaining)
        bb.get(dest, position, remaining)
        dest
      }
    }
  }

} 
Example 144
Source File: ShapelessPicklers.scala    From boopickle   with Apache License 2.0 5 votes vote down vote up
package boopickle.shapeless

import boopickle._
import _root_.shapeless._

import java.nio.ByteBuffer
import scala.reflect.ClassTag

trait ShapelessPicklers extends TransformPicklers {
  implicit def hconsPickler[H, T <: HList](implicit hp: Lazy[Pickler[H]], tp: Lazy[Pickler[T]]): Pickler[H :: T] =
    new Pickler[H :: T] {
      override def pickle(list: H :: T)(implicit state: PickleState): Unit = {
        val head :: tail = list
        hp.value.pickle(head)
        tp.value.pickle(tail)
      }

      override def unpickle(implicit state: UnpickleState): H :: T = {
        val head = hp.value.unpickle
        val tail = tp.value.unpickle
        head :: tail
      }
    }

  implicit val hnilPickler: Pickler[HNil] = new Pickler[HNil] {
    override def pickle(list: HNil)(implicit state: PickleState): Unit = ()
    override def unpickle(implicit state: UnpickleState): HNil         = HNil
  }

  implicit def genericPickler[A, B](implicit gen: Generic.Aux[A, B], rp: Lazy[Pickler[B]]): Pickler[A] = new Pickler[A] {
    override def pickle(list: A)(implicit state: PickleState): Unit = rp.value.pickle(gen.to(list))
    override def unpickle(implicit state: UnpickleState): A         = gen.from(rp.value.unpickle)
  }

  implicit def coproductInlPickler[A, B <: Coproduct](implicit ap: Lazy[Pickler[A]]): Pickler[Inl[A, B]] = {
    transformPickler[Inl[A, B], A](a => Inl(a))(_.head)(ap.value)
  }

  implicit def coproductInrPickler[A, B <: Coproduct](implicit bp: Lazy[Pickler[B]]): Pickler[Inr[A, B]] = {
    transformPickler[Inr[A, B], B](a => Inr(a))(_.tail)(bp.value)
  }

  implicit def coproductPickler[A: ClassTag, B <: Coproduct: ClassTag](implicit ap: Lazy[Pickler[A]],
                                                                       bp: Lazy[Pickler[B]]): Pickler[A :+: B] = {
    CompositePickler[A :+: B].addConcreteType[Inl[A, B]].addConcreteType[Inr[A, B]]
  }

  implicit def cnilPickler: Pickler[CNil] = new Pickler[CNil] {
    override def pickle(list: CNil)(implicit state: PickleState): Unit = ()
    override def unpickle(implicit state: UnpickleState): CNil         = ??? // CNil should never be reached
  }
}
object ShapelessPicklers extends ShapelessPicklers 
Example 145
Source File: StringCodec.scala    From boopickle   with Apache License 2.0 5 votes vote down vote up
package boopickle

import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets

object StringCodec extends StringCodecBase {
  override def decodeUTF8(len: Int, buf: ByteBuffer): String = {
    val a = new Array[Byte](len)
    buf.get(a)
    new String(a, StandardCharsets.UTF_8)
  }

  override def encodeUTF8(str: String): ByteBuffer = {
    ByteBuffer.wrap(str.getBytes(StandardCharsets.UTF_8))
  }

  override def decodeUTF16(len: Int, buf: ByteBuffer): String = {
    val a = new Array[Byte](len)
    buf.get(a)
    new String(a, StandardCharsets.UTF_16LE)
  }

  override def encodeUTF16(str: String): ByteBuffer = {
    ByteBuffer.wrap(str.getBytes(StandardCharsets.UTF_16LE))
  }
} 
Example 146
Source File: BufferPoolTestsJVM.scala    From boopickle   with Apache License 2.0 5 votes vote down vote up
package boopickle

import java.nio.{ByteBuffer, ByteOrder}

import utest._

object BufferPoolTestsJVM extends TestSuite {

  override def tests = Tests {
    "MultiThread" - {
      val pool  = BufferPool
      val count = 100000
      def runner = new Runnable {
        override def run(): Unit = {
          var i = 0
          while (i < count) {
            val bb1 = pool
              .allocate(ByteBufferProvider.initSize)
              .getOrElse(ByteBuffer.allocate(ByteBufferProvider.initSize))
              .order(ByteOrder.LITTLE_ENDIAN)
            val bb2 = pool
              .allocate(ByteBufferProvider.expandSize)
              .getOrElse(ByteBuffer.allocate(ByteBufferProvider.expandSize))
              .order(ByteOrder.LITTLE_ENDIAN)
            pool.release(bb1)
            pool.release(bb2)
            pool.release(ByteBuffer.allocate(ByteBufferProvider.initSize).order(ByteOrder.LITTLE_ENDIAN))
            i += 1
          }
        }
      }
      // warmup
      runner.run()
      runner.run()
      System.gc()
      // run in a single thread
      var startTime = System.nanoTime()
      runner.run()
      var endTime = System.nanoTime()
      println(s"Single thread: ${(endTime - startTime) / 1000}")
      var t1 = new Thread(runner)
      var t2 = new Thread(runner)
      startTime = System.nanoTime()
      t1.start()
      t2.start()
      t1.join()
      t2.join()
      endTime = System.nanoTime()
      println(s"Two threads: ${(endTime - startTime) / 1000}")
      startTime = System.nanoTime()
      t1 = new Thread(runner)
      t2 = new Thread(runner)
      val t3 = new Thread(runner)
      t1.start()
      t2.start()
      t3.start()
      t1.join()
      t2.join()
      t3.join()
      endTime = System.nanoTime()
      println(s"Three threads: ${(endTime - startTime) / 1000}")
    }
  }
} 
Example 147
Source File: PerfTests.scala    From boopickle   with Apache License 2.0 5 votes vote down vote up
package external

import java.nio.ByteBuffer

import boopickle.{DecoderSize, EncoderSize, PickleState, UnpickleState}
import boopickle.Default._
import utest._

object PerfTests extends TestSuite {
  def tests = Tests {
    "Performance" - {
      case class Test(i: Int, s: String)
      // generate data
      val template = (0 until 500).map(i => Test(i, (i / 2).toString * 20))
      val data     = (0 until 2000).map(i => template(i % template.size))
      def dedupTest(topic: String, pState: => PickleState, uState: ByteBuffer => UnpickleState) = {
        def testRun = {
          val pickleState   = pState
          val start         = System.nanoTime()
          val bb            = pickleState.pickle(data).toByteBuffer
          val middle        = System.nanoTime()
          val unpickleState = uState(bb)
          val uData         = unpickleState.unpickle[Seq[Test]]
          val end           = System.nanoTime()
          (middle - start, end - middle, uData, bb)
        }
        val (eTime, dTime, uData, bb) =
          (0 until 3).foldLeft((Long.MaxValue, Long.MaxValue, Seq.empty[Test], null.asInstanceOf[ByteBuffer])) {
            case ((enc, dec, _, _), idx) =>
              val (e, d, data, bb) = testRun
              (e min enc, d min dec, data, bb)
          }
        println(s"$topic -- Pickle time: ${eTime / 1000}, Unpickle time:  ${dTime / 1000}")
        println(s"Data size ${bb.capacity()}")
        assert(uData == data)
      }
      "Deduplication" - {
        dedupTest("With dedup",
                  new PickleState(new EncoderSize, true, true),
                  bb => new UnpickleState(new DecoderSize(bb), true, true))
      }
      "NoDeduplication" - {
        dedupTest("Without dedup",
                  new PickleState(new EncoderSize, false, false),
                  bb => new UnpickleState(new DecoderSize(bb), false, false))
      }
    }
  }
} 
Example 148
Source File: BufferProviderTests.scala    From boopickle   with Apache License 2.0 5 votes vote down vote up
package boopickle

import java.nio.ByteBuffer

import scala.util.Random

import external.Banana
import utest._
import boopickle.Default._

object BufferProviderTests extends TestSuite {

  override def tests = Tests {
    "asByteBuffersProperOrder" - {

      val input: Seq[Banana] = Iterator.tabulate(1000)(_ => Banana(Random.nextDouble)).toVector
      val bbs                = Pickle.intoByteBuffers(input)
      assert(bbs.size > 1)

      val mergedBb = ByteBuffer.allocate(bbs.map(_.remaining).sum)
      bbs.foreach(mergedBb.put)
      mergedBb.flip()

      val output = Unpickle[Seq[Banana]].fromBytes(mergedBb)
      assert(output == input)
    }
  }
} 
Example 149
Source File: StringCodecTests.scala    From boopickle   with Apache License 2.0 5 votes vote down vote up
package boopickle

import java.nio.ByteBuffer

import utest._

import boopickle.DefaultBasic._

object StringCodecTests extends TestSuite {

  override def tests = Tests {
    'LargeString {
      // test data
      val data = new String(Array.fill[Byte](400000)('A'))
      // create encoded strings of various lengths
      def createBB(size: Int) = Pickle.intoBytes(data.substring(0, size))

      val sizes = Seq(1000, 4095, 4096, 4097, 8192, 70000, 200000, 400000)
      val bufs  = sizes.map(createBB)

      val strings = bufs.map(b => Unpickle[String].fromBytes(b))
      sizes.zip(strings).foreach {
        case (size, str) => assert(str.length == size)
      }
    }
  }
} 
Example 150
Source File: StringCodec.scala    From boopickle   with Apache License 2.0 5 votes vote down vote up
package boopickle

import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets

object StringCodec extends StringCodecBase {
  override def decodeUTF8(len: Int, buf: ByteBuffer): String = {
    val a = new Array[Byte](len)
    buf.get(a)
    new String(a, StandardCharsets.UTF_8)
  }

  override def encodeUTF8(str: String): ByteBuffer = {
    ByteBuffer.wrap(str.getBytes(StandardCharsets.UTF_8))
  }

  override def decodeUTF16(len: Int, buf: ByteBuffer): String = {
    val a = new Array[Byte](len)
    buf.get(a)
    new String(a, StandardCharsets.UTF_16LE)
  }

  override def encodeUTF16(str: String): ByteBuffer = {
    ByteBuffer.wrap(str.getBytes(StandardCharsets.UTF_16LE))
  }
} 
Example 151
Source File: Show.scala    From boopickle   with Apache License 2.0 5 votes vote down vote up
import java.nio.ByteBuffer

import scalatags.JsDom.all._

trait Show[A] {
  def toHTML(a: A): Modifier
}

private def printByteBuffer(bb: ByteBuffer): List[String] = {
  val data = Array.ofDim[Byte](bb.remaining())
  bb.duplicate().get(data)
  data.grouped(16).toList.map { d =>
    val hex = d.map(c => "%02X " format (c & 0xFF)).mkString
    val str = d
      .collect {
        case ascii if ascii >= 0x20 && ascii < 0x80 => ascii
        case _                                      => '.'.toByte
      }
      .map(_.toChar)
      .mkString
    hex.padTo(16 * 3, ' ') + str
  }
}

implicit val showByteBuffer: Show[ByteBuffer] = new Show[ByteBuffer] {
  def toHTML(a: ByteBuffer): Modifier =
    pre(
      style := "background-color: #eee; padding: 10px;",
      printByteBuffer(a).mkString("\n")
    )
}

def show[A](a: A)(implicit s: Show[A]) = Fiddle.print(s.toHTML(a)) 
Example 152
Source File: AsyncHttpClientScalazBackend.scala    From sttp   with Apache License 2.0 5 votes vote down vote up
package sttp.client.asynchttpclient.scalaz

import java.nio.ByteBuffer

import io.netty.buffer.ByteBuf
import org.asynchttpclient.{
  AsyncHttpClient,
  AsyncHttpClientConfig,
  BoundRequestBuilder,
  DefaultAsyncHttpClient,
  DefaultAsyncHttpClientConfig
}
import org.reactivestreams.Publisher
import scalaz.concurrent.Task
import sttp.client.asynchttpclient.{AsyncHttpClientBackend, WebSocketHandler}
import sttp.client.impl.scalaz.TaskMonadAsyncError
import sttp.client.testing.SttpBackendStub
import sttp.client.{FollowRedirectsBackend, SttpBackend, SttpBackendOptions}

class AsyncHttpClientScalazBackend private (
    asyncHttpClient: AsyncHttpClient,
    closeClient: Boolean,
    customizeRequest: BoundRequestBuilder => BoundRequestBuilder
) extends AsyncHttpClientBackend[Task, Nothing](asyncHttpClient, TaskMonadAsyncError, closeClient, customizeRequest) {
  override protected def streamBodyToPublisher(s: Nothing): Publisher[ByteBuf] =
    s // nothing is everything

  override protected def publisherToStreamBody(p: Publisher[ByteBuffer]): Nothing =
    throw new IllegalStateException("This backend does not support streaming")
}

object AsyncHttpClientScalazBackend {
  private def apply(
      asyncHttpClient: AsyncHttpClient,
      closeClient: Boolean,
      customizeRequest: BoundRequestBuilder => BoundRequestBuilder
  ): SttpBackend[Task, Nothing, WebSocketHandler] =
    new FollowRedirectsBackend[Task, Nothing, WebSocketHandler](
      new AsyncHttpClientScalazBackend(asyncHttpClient, closeClient, customizeRequest)
    )

  def apply(
      options: SttpBackendOptions = SttpBackendOptions.Default,
      customizeRequest: BoundRequestBuilder => BoundRequestBuilder = identity
  ): Task[SttpBackend[Task, Nothing, WebSocketHandler]] =
    Task.delay(
      AsyncHttpClientScalazBackend(AsyncHttpClientBackend.defaultClient(options), closeClient = true, customizeRequest)
    )

  def usingConfig(
      cfg: AsyncHttpClientConfig,
      customizeRequest: BoundRequestBuilder => BoundRequestBuilder = identity
  ): Task[SttpBackend[Task, Nothing, WebSocketHandler]] =
    Task.delay(AsyncHttpClientScalazBackend(new DefaultAsyncHttpClient(cfg), closeClient = true, customizeRequest))

  
  def stub: SttpBackendStub[Task, Nothing, WebSocketHandler] = SttpBackendStub(TaskMonadAsyncError)
} 
Example 153
Source File: AsyncHttpClientFs2Backend.scala    From sttp   with Apache License 2.0 5 votes vote down vote up
package sttp.client.asynchttpclient.fs2

import java.io.File
import java.nio.ByteBuffer

import cats.effect._
import cats.effect.implicits._
import cats.implicits._
import fs2.{Chunk, Stream}
import fs2.interop.reactivestreams._
import io.netty.buffer.{ByteBuf, Unpooled}
import org.asynchttpclient.{Request => _, Response => _, _}
import org.reactivestreams.Publisher
import sttp.client.asynchttpclient.{AsyncHttpClientBackend, WebSocketHandler}
import sttp.client.impl.cats.CatsMonadAsyncError
import sttp.client.internal._
import sttp.client.testing.SttpBackendStub
import sttp.client.ws.WebSocketResponse
import sttp.client.{FollowRedirectsBackend, SttpBackend, SttpBackendOptions, _}

import scala.concurrent.ExecutionContext
import scala.language.higherKinds

class AsyncHttpClientFs2Backend[F[_]: ConcurrentEffect: ContextShift] private (
    asyncHttpClient: AsyncHttpClient,
    closeClient: Boolean,
    customizeRequest: BoundRequestBuilder => BoundRequestBuilder
) extends AsyncHttpClientBackend[F, Stream[F, Byte]](
      asyncHttpClient,
      new CatsMonadAsyncError,
      closeClient,
      customizeRequest
    ) {
  override def send[T](r: Request[T, Stream[F, Byte]]): F[Response[T]] = {
    super.send(r).guarantee(implicitly[ContextShift[F]].shift)
  }

  override def openWebsocket[T, WS_RESULT](
      r: Request[T, Stream[F, Byte]],
      handler: WebSocketHandler[WS_RESULT]
  ): F[WebSocketResponse[WS_RESULT]] = super.openWebsocket(r, handler).guarantee(ContextShift[F].shift)

  override protected def streamBodyToPublisher(s: Stream[F, Byte]): Publisher[ByteBuf] =
    s.chunks.map(c => Unpooled.wrappedBuffer(c.toArray)).toUnicastPublisher

  override protected def publisherToStreamBody(p: Publisher[ByteBuffer]): Stream[F, Byte] =
    p.toStream[F].flatMap(buf => Stream.chunk(Chunk.byteBuffer(buf)))

  override protected def publisherToBytes(p: Publisher[ByteBuffer]): F[Array[Byte]] = {
    p.toStream[F]
      .compile
      .fold(ByteBuffer.allocate(0))(concatByteBuffers)
      .map(_.array())
  }

  override protected def publisherToFile(p: Publisher[ByteBuffer], f: File): F[Unit] = {
    p.toStream[F]
      .flatMap(b => Stream.emits(b.array()))
      .through(fs2.io.file.writeAll(f.toPath, Blocker.liftExecutionContext(ExecutionContext.global)))
      .compile
      .drain
  }
}

object AsyncHttpClientFs2Backend {
  private def apply[F[_]: ConcurrentEffect: ContextShift](
      asyncHttpClient: AsyncHttpClient,
      closeClient: Boolean,
      customizeRequest: BoundRequestBuilder => BoundRequestBuilder
  ): SttpBackend[F, Stream[F, Byte], WebSocketHandler] =
    new FollowRedirectsBackend(new AsyncHttpClientFs2Backend(asyncHttpClient, closeClient, customizeRequest))

  def apply[F[_]: ConcurrentEffect: ContextShift](
      options: SttpBackendOptions = SttpBackendOptions.Default,
      customizeRequest: BoundRequestBuilder => BoundRequestBuilder = identity
  ): F[SttpBackend[F, Stream[F, Byte], WebSocketHandler]] =
    implicitly[Sync[F]]
      .delay(apply[F](AsyncHttpClientBackend.defaultClient(options), closeClient = true, customizeRequest))

  
  def stub[F[_]: Concurrent]: SttpBackendStub[F, Stream[F, ByteBuffer], WebSocketHandler] =
    SttpBackendStub(new CatsMonadAsyncError())
} 
Example 154
Source File: AsyncHttpClientFutureBackend.scala    From sttp   with Apache License 2.0 5 votes vote down vote up
package sttp.client.asynchttpclient.future

import java.nio.ByteBuffer

import io.netty.buffer.ByteBuf
import org.asynchttpclient.{
  AsyncHttpClient,
  AsyncHttpClientConfig,
  BoundRequestBuilder,
  DefaultAsyncHttpClient,
  DefaultAsyncHttpClientConfig
}
import org.reactivestreams.Publisher
import sttp.client.asynchttpclient.{AsyncHttpClientBackend, WebSocketHandler}
import sttp.client.monad.FutureMonad
import sttp.client.testing.SttpBackendStub
import sttp.client.{FollowRedirectsBackend, SttpBackend, SttpBackendOptions}

import scala.concurrent.{ExecutionContext, Future}

class AsyncHttpClientFutureBackend private (
    asyncHttpClient: AsyncHttpClient,
    closeClient: Boolean,
    customizeRequest: BoundRequestBuilder => BoundRequestBuilder
)(implicit
    ec: ExecutionContext
) extends AsyncHttpClientBackend[Future, Nothing](asyncHttpClient, new FutureMonad, closeClient, customizeRequest) {
  override protected def streamBodyToPublisher(s: Nothing): Publisher[ByteBuf] =
    s // nothing is everything

  override protected def publisherToStreamBody(p: Publisher[ByteBuffer]): Nothing =
    throw new IllegalStateException("This backend does not support streaming")
}

object AsyncHttpClientFutureBackend {
  private def apply(
      asyncHttpClient: AsyncHttpClient,
      closeClient: Boolean,
      customizeRequest: BoundRequestBuilder => BoundRequestBuilder
  )(implicit
      ec: ExecutionContext
  ): SttpBackend[Future, Nothing, WebSocketHandler] =
    new FollowRedirectsBackend[Future, Nothing, WebSocketHandler](
      new AsyncHttpClientFutureBackend(asyncHttpClient, closeClient, customizeRequest)
    )

  
  def stub(implicit
      ec: ExecutionContext = ExecutionContext.global
  ): SttpBackendStub[Future, Nothing, WebSocketHandler] =
    SttpBackendStub(new FutureMonad())
} 
Example 155
Source File: DelegatingWebSocketListener.scala    From sttp   with Apache License 2.0 5 votes vote down vote up
package sttp.client.httpclient

import java.net.http.WebSocket
import java.net.http.WebSocket.Listener
import java.nio.ByteBuffer
import java.util.concurrent.CompletionStage
import java.util.concurrent.atomic.AtomicBoolean

private[httpclient] class DelegatingWebSocketListener[WS_RESULT](
    delegate: Listener,
    onInitialOpen: WebSocket => Unit,
    onInitialError: Throwable => Unit
) extends Listener {
  private val initialised = new AtomicBoolean(false)

  override def onOpen(webSocket: WebSocket): Unit = {
    if (!initialised.getAndSet(true)) {
      onInitialOpen(webSocket)
    }
    delegate.onOpen(webSocket)
  }

  override def onText(webSocket: WebSocket, data: CharSequence, last: Boolean): CompletionStage[_] = {
    delegate.onText(webSocket, data, last)
  }

  override def onBinary(webSocket: WebSocket, data: ByteBuffer, last: Boolean): CompletionStage[_] = {
    delegate.onBinary(webSocket, data, last)
  }

  override def onPing(webSocket: WebSocket, message: ByteBuffer): CompletionStage[_] = {
    delegate.onPing(webSocket, message)
  }

  override def onPong(webSocket: WebSocket, message: ByteBuffer): CompletionStage[_] = {
    delegate.onPong(webSocket, message)
  }

  override def onClose(webSocket: WebSocket, statusCode: Int, reason: String): CompletionStage[_] = {
    delegate.onClose(webSocket, statusCode, reason)
  }
  override def onError(webSocket: WebSocket, error: Throwable): Unit = {
    if (!initialised.getAndSet(true)) {
      onInitialError(error)
    }
    delegate.onError(webSocket, error)
  }
} 
Example 156
Source File: HttpClientMonixBackend.scala    From sttp   with Apache License 2.0 5 votes vote down vote up
package sttp.client.httpclient.monix

import java.io.InputStream
import java.net.http.HttpRequest.BodyPublishers
import java.net.http.{HttpClient, HttpRequest}
import java.nio.ByteBuffer

import cats.effect.Resource
import monix.eval.Task
import monix.execution.Scheduler
import monix.reactive.Observable
import org.reactivestreams.FlowAdapters
import sttp.client.httpclient.HttpClientBackend.EncodingHandler
import sttp.client.httpclient.{HttpClientAsyncBackend, HttpClientBackend, WebSocketHandler}
import sttp.client.impl.monix.TaskMonadAsyncError
import sttp.client.testing.SttpBackendStub
import sttp.client.{FollowRedirectsBackend, SttpBackend, SttpBackendOptions}

import scala.util.{Success, Try}

class HttpClientMonixBackend private (
    client: HttpClient,
    closeClient: Boolean,
    customizeRequest: HttpRequest => HttpRequest,
    customEncodingHandler: EncodingHandler
)(implicit s: Scheduler)
    extends HttpClientAsyncBackend[Task, Observable[ByteBuffer]](
      client,
      TaskMonadAsyncError,
      closeClient,
      customizeRequest,
      customEncodingHandler
    ) {
  override def streamToRequestBody(stream: Observable[ByteBuffer]): Task[HttpRequest.BodyPublisher] = {
    monad.eval(BodyPublishers.fromPublisher(FlowAdapters.toFlowPublisher(stream.toReactivePublisher)))
  }

  override def responseBodyToStream(responseBody: InputStream): Try[Observable[ByteBuffer]] = {
    Success(
      Observable
        .fromInputStream(Task.now(responseBody))
        .map(ByteBuffer.wrap)
        .guaranteeCase(_ => Task(responseBody.close()))
    )
  }
}

object HttpClientMonixBackend {
  private def apply(
      client: HttpClient,
      closeClient: Boolean,
      customizeRequest: HttpRequest => HttpRequest,
      customEncodingHandler: EncodingHandler
  )(implicit
      s: Scheduler
  ): SttpBackend[Task, Observable[ByteBuffer], WebSocketHandler] =
    new FollowRedirectsBackend(
      new HttpClientMonixBackend(client, closeClient, customizeRequest, customEncodingHandler)(s)
    )

  def apply(
      options: SttpBackendOptions = SttpBackendOptions.Default,
      customizeRequest: HttpRequest => HttpRequest = identity,
      customEncodingHandler: EncodingHandler = PartialFunction.empty
  )(implicit
      s: Scheduler = Scheduler.global
  ): Task[SttpBackend[Task, Observable[ByteBuffer], WebSocketHandler]] =
    Task.eval(
      HttpClientMonixBackend(
        HttpClientBackend.defaultClient(options),
        closeClient = true,
        customizeRequest,
        customEncodingHandler
      )(s)
    )

  def resource(
      options: SttpBackendOptions = SttpBackendOptions.Default,
      customizeRequest: HttpRequest => HttpRequest = identity,
      customEncodingHandler: EncodingHandler = PartialFunction.empty
  )(implicit
      s: Scheduler = Scheduler.global
  ): Resource[Task, SttpBackend[Task, Observable[ByteBuffer], WebSocketHandler]] =
    Resource.make(apply(options, customizeRequest, customEncodingHandler))(_.close())

  def usingClient(
      client: HttpClient,
      customizeRequest: HttpRequest => HttpRequest = identity,
      customEncodingHandler: EncodingHandler = PartialFunction.empty
  )(implicit s: Scheduler = Scheduler.global): SttpBackend[Task, Observable[ByteBuffer], WebSocketHandler] =
    HttpClientMonixBackend(client, closeClient = false, customizeRequest, customEncodingHandler)(s)

  
  def stub: SttpBackendStub[Task, Observable[ByteBuffer], WebSocketHandler] = SttpBackendStub(TaskMonadAsyncError)
} 
Example 157
Source File: HttpClientHighLevelMonixWebsocketTest.scala    From sttp   with Apache License 2.0 5 votes vote down vote up
package sttp.client.httpclient.monix

import java.nio.ByteBuffer

import monix.eval.Task
import monix.execution.Scheduler.Implicits.global
import monix.reactive.Observable
import sttp.client._
import sttp.client.httpclient.WebSocketHandler
import sttp.client.impl.monix.{TaskMonadAsyncError, convertMonixTaskToFuture}
import sttp.client.monad.MonadError
import sttp.client.testing.ConvertToFuture
import sttp.client.testing.websocket.HighLevelWebsocketTest
import sttp.client.ws.WebSocket
import sttp.client.testing.HttpTest.wsEndpoint

import scala.concurrent.duration._

class HttpClientHighLevelMonixWebsocketTest extends HighLevelWebsocketTest[Task, WebSocketHandler] {
  implicit val backend: SttpBackend[Task, Observable[ByteBuffer], WebSocketHandler] =
    HttpClientMonixBackend().runSyncUnsafe()
  implicit val convertToFuture: ConvertToFuture[Task] = convertMonixTaskToFuture
  implicit val monad: MonadError[Task] = TaskMonadAsyncError

  override def createHandler: Option[Int] => Task[WebSocketHandler[WebSocket[Task]]] = _ => MonixWebSocketHandler()

  it should "handle backpressure correctly" in {
    basicRequest
      .get(uri"$wsEndpoint/ws/echo")
      .openWebsocketF(createHandler(None))
      .flatMap { response =>
        val ws = response.result
        send(ws, 1000) >> eventually(10.millis, 500) { ws.isOpen.map(_ shouldBe true) }
      }
      .toFuture()
  }

  override def eventually[T](interval: FiniteDuration, attempts: Int)(f: => Task[T]): Task[T] = {
    (Task.sleep(interval) >> f).onErrorRestart(attempts.toLong)
  }
} 
Example 158
Source File: MonixStreamingTest.scala    From sttp   with Apache License 2.0 5 votes vote down vote up
package sttp.client.impl.monix

import java.nio.ByteBuffer

import monix.eval.Task
import monix.reactive.Observable
import sttp.client.testing.ConvertToFuture
import sttp.client.testing.streaming.StreamingTest

abstract class MonixStreamingTest extends StreamingTest[Task, Observable[ByteBuffer]] {

  override implicit val convertToFuture: ConvertToFuture[Task] = convertMonixTaskToFuture

  override def bodyProducer(chunks: Iterable[Array[Byte]]): Observable[ByteBuffer] =
    Observable
      .fromIterable(chunks)
      .map(ByteBuffer.wrap)

  override def bodyConsumer(stream: Observable[ByteBuffer]): Task[String] =
    stream
      .flatMap(v => Observable.fromIterable(v.array()))
      .toListL
      .map(bs => new String(bs.toArray, "utf8"))
} 
Example 159
Source File: package.scala    From sttp   with Apache License 2.0 5 votes vote down vote up
package sttp.client

import java.io.{ByteArrayOutputStream, InputStream, OutputStream}
import java.nio.{Buffer, ByteBuffer}

import scala.annotation.{implicitNotFound, tailrec}

package object internal {
  private[client] def contentTypeWithCharset(ct: String, charset: String): String =
    s"$ct; charset=$charset"

  private[client] def charsetFromContentType(ct: String): Option[String] =
    ct.split(";").map(_.trim.toLowerCase).collectFirst {
      case s if s.startsWith("charset=") && s.substring(8).trim != "" => s.substring(8).trim
    }

  private[client] def transfer(is: InputStream, os: OutputStream): Unit = {
    var read = 0
    val buf = new Array[Byte](1024)

    @tailrec
    def transfer(): Unit = {
      read = is.read(buf, 0, buf.length)
      if (read != -1) {
        os.write(buf, 0, read)
        transfer()
      }
    }

    transfer()
  }

  private[client] def toByteArray(is: InputStream): Array[Byte] = {
    val os = new ByteArrayOutputStream
    transfer(is, os)
    os.toByteArray
  }

  private[client] def concatByteBuffers(bb1: ByteBuffer, bb2: ByteBuffer): ByteBuffer = {
    val buf = ByteBuffer
      .allocate(bb1.array().length + bb2.array().length)
      .put(bb1)
      .put(bb2)
    // rewind() returns Buffer in Java8, and ByteBuffer in Java11
    // calling the method from the base class to avoid NoSuchMethodError
    (buf: Buffer).rewind()
    buf
  }

  
  private[client] def sanitizeCharset(charset: String): String = {
    val c2 = charset.trim()
    val c3 = if (c2.startsWith("\"")) c2.substring(1) else c2
    if (c3.endsWith("\"")) c3.substring(0, c3.length - 1) else c3
  }

  @implicitNotFound(
    "This is a partial request, the method & url are not specified. Use " +
      ".get(...), .post(...) etc. to obtain a non-partial request."
  )
  private[client] type IsIdInRequest[U[_]] = U[Unit] =:= Identity[Unit]

  private[client] val Utf8 = "utf-8"
  private[client] val Iso88591 = "iso-8859-1"
  private[client] val CrLf = "\r\n"
} 
Example 160
Source File: RequestBody.scala    From sttp   with Apache License 2.0 5 votes vote down vote up
package sttp.client

import java.io.InputStream
import java.nio.ByteBuffer

import sttp.model._
import sttp.client.internal.SttpFile
import sttp.model.internal.UriCompatibility

import scala.collection.immutable.Seq

sealed trait RequestBody[+S]
case object NoBody extends RequestBody[Nothing]

sealed trait BasicRequestBody extends RequestBody[Nothing] {
  def defaultContentType: Option[MediaType]
}

case class StringBody(
    s: String,
    encoding: String,
    defaultContentType: Option[MediaType] = Some(MediaType.TextPlain)
) extends BasicRequestBody

case class ByteArrayBody(
    b: Array[Byte],
    defaultContentType: Option[MediaType] = Some(MediaType.ApplicationOctetStream)
) extends BasicRequestBody

case class ByteBufferBody(
    b: ByteBuffer,
    defaultContentType: Option[MediaType] = Some(MediaType.ApplicationOctetStream)
) extends BasicRequestBody

case class InputStreamBody(
    b: InputStream,
    defaultContentType: Option[MediaType] = Some(MediaType.ApplicationOctetStream)
) extends BasicRequestBody

case class FileBody(
    f: SttpFile,
    defaultContentType: Option[MediaType] = Some(MediaType.ApplicationOctetStream)
) extends BasicRequestBody

case class StreamBody[S](s: S) extends RequestBody[S]

case class MultipartBody(parts: Seq[Part[BasicRequestBody]]) extends RequestBody[Nothing]

object RequestBody {
  private[client] def paramsToStringBody(fs: Seq[(String, String)], encoding: String): StringBody = {
    val b = fs
      .map {
        case (key, value) =>
          UriCompatibility.encodeQuery(key, encoding) + "=" +
            UriCompatibility.encodeQuery(value, encoding)
      }
      .mkString("&")

    StringBody(b, encoding)
  }
} 
Example 161
Source File: ByteStringBytes.scala    From swave   with Mozilla Public License 2.0 5 votes vote down vote up
package swave.compat.akka.impl

import java.io.OutputStream
import java.nio.charset.{CharacterCodingException, Charset}
import java.nio.{ByteBuffer, CharBuffer}
import java.util
import scala.annotation.tailrec
import scala.collection.GenTraversableOnce
import akka.util.ByteString
import swave.core.io.Bytes
import swave.core.macros._

class ByteStringBytes extends Bytes[ByteString] {

  ///////////////// CONSTRUCTION ///////////////////

  def empty = ByteString.empty
  def fill[A: Integral](size: Long)(byte: A) = {
    requireArg(0 <= size && size <= Int.MaxValue, "`size` must be >= 0 and <= Int.MaxValue")
    val b = implicitly[Integral[A]].toInt(byte).toByte
    apply(Array.fill(size.toInt)(b))
  }
  def apply(array: Array[Byte]) = ByteString(array)
  def apply(bytes: Array[Byte], offset: Int, length: Int) =
    ByteString(util.Arrays.copyOfRange(bytes, offset, offset + length))
  def apply[A: Integral](bytes: A*) = {
    val integral = implicitly[Integral[A]]
    val buf      = new Array[Byte](bytes.size)
    @tailrec def rec(ix: Int): ByteString =
      if (ix < buf.length) {
        buf(ix) = integral.toInt(bytes(ix)).toByte
        rec(ix + 1)
      } else view(buf)
    rec(0)
  }
  def apply(bytes: Vector[Byte])                  = ByteString(bytes: _*)
  def apply(buffer: ByteBuffer)                   = ByteString(buffer)
  def apply(bs: GenTraversableOnce[Byte])         = ByteString(bs.toArray)
  def view(bytes: Array[Byte])                    = ByteString(bytes) // no view-like constructor available on ByteStrings
  def view(bytes: ByteBuffer)                     = ByteString(bytes) // no view-like constructor available on ByteStrings
  def encodeString(str: String, charset: Charset) = ByteString(str, charset.name)
  def encodeStringStrict(str: String, charset: Charset) =
    try Right(ByteString(charset.newEncoder.encode(CharBuffer.wrap(str))))
    catch { case e: CharacterCodingException ⇒ Left(e) }

  ///////////////// QUERY ///////////////////

  def size(value: ByteString): Long = value.size.toLong
  def byteAt(value: ByteString, ix: Long): Byte = {
    requireArg(0 <= ix && ix <= Int.MaxValue, "`ix` must be >= 0 and <= Int.MaxValue")
    value(ix.toInt)
  }
  def indexOfSlice(value: ByteString, slice: ByteString, startIx: Long): Long = {
    requireArg(0 <= startIx && startIx <= Int.MaxValue, "`startIx` must be >= 0 and <= Int.MaxValue")
    value.indexOfSlice(slice, startIx.toInt).toLong
  }

  ///////////////// TRANSFORMATION TO ByteString ///////////////////

  def update(value: ByteString, ix: Long, byte: Byte) = concat(concat(take(value, ix), byte), drop(value, ix + 1))
  def concat(value: ByteString, other: ByteString)    = value ++ other
  def concat(value: ByteString, byte: Byte)           = value ++ ByteString(byte)
  def concat(byte: Byte, value: ByteString)           = ByteString(byte) ++ value
  def drop(value: ByteString, n: Long) = {
    requireArg(0 <= n && n <= Int.MaxValue, "`n` must be >= 0 and <= Int.MaxValue")
    value.drop(n.toInt)
  }
  def take(value: ByteString, n: Long) = {
    requireArg(0 <= n && n <= Int.MaxValue, "`n` must be >= 0 and <= Int.MaxValue")
    value.take(n.toInt)
  }
  def map(value: ByteString, f: Byte ⇒ Byte) = value.map(f)
  def reverse(value: ByteString)             = value.reverse
  def compact(value: ByteString)             = value.compact

  ///////////////// TRANSFORMATION TO OTHER TYPES ///////////////////

  def toArray(value: ByteString)                                   = value.toArray
  def copyToArray(value: ByteString, xs: Array[Byte], offset: Int) = value.copyToArray(xs, offset)
  def copyToArray(value: ByteString, sourceOffset: Long, xs: Array[Byte], destOffset: Int, len: Int) =
    drop(value, sourceOffset).copyToArray(xs, destOffset, len)
  def copyToBuffer(value: ByteString, buffer: ByteBuffer): Int = value.copyToBuffer(buffer)
  def copyToOutputStream(value: ByteString, s: OutputStream) = {
    @tailrec def rec(ix: Int, size: Int): Unit = if (ix < size) { s.write(value(ix).toInt); rec(ix + 1, size) }
    rec(0, value.size)
  }
  def toByteBuffer(value: ByteString)                   = value.toByteBuffer
  def toIndexedSeq(value: ByteString): IndexedSeq[Byte] = value
  def toSeq(value: ByteString): Seq[Byte]               = value
  def decodeString(value: ByteString, charset: Charset): Either[CharacterCodingException, String] =
    try Right(charset.newDecoder.decode(toByteBuffer(value)).toString)
    catch { case e: CharacterCodingException ⇒ Left(e) }

  ///////////////// ITERATION ///////////////////

  def foldLeft[A](value: ByteString, z: A, f: (A, Byte) ⇒ A)  = value.foldLeft(z)(f)
  def foldRight[A](value: ByteString, z: A, f: (Byte, A) ⇒ A) = value.foldRight(z)(f)
  def foreach(value: ByteString, f: Byte ⇒ Unit)              = value.foreach(f)

} 
Example 162
Source File: ByteVectorBytes.scala    From swave   with Mozilla Public License 2.0 5 votes vote down vote up
package swave.compat.scodec.impl

import java.io.OutputStream
import java.nio.ByteBuffer
import java.nio.charset.{CharacterCodingException, Charset}
import scala.collection.GenTraversableOnce
import scodec.bits.ByteVector
import swave.core.io.Bytes

class ByteVectorBytes extends Bytes[ByteVector] {

  ///////////////// CONSTRUCTION ///////////////////

  def empty                                               = ByteVector.empty
  def fill[A: Integral](size: Long)(byte: A)              = ByteVector.fill(size)(byte)
  def apply(array: Array[Byte])                           = ByteVector(array)
  def apply(bytes: Array[Byte], offset: Int, length: Int) = ByteVector(bytes, offset, length)
  def apply[A: Integral](bytes: A*)                       = ByteVector(bytes: _*)
  def apply(bytes: Vector[Byte])                          = ByteVector(bytes)
  def apply(buffer: ByteBuffer)                           = ByteVector(buffer)
  def apply(bs: GenTraversableOnce[Byte])                 = ByteVector(bs)
  def view(bytes: Array[Byte])                            = ByteVector(bytes)
  def view(bytes: ByteBuffer)                             = ByteVector(bytes)
  def encodeString(str: String, charset: Charset)         = if (str.isEmpty) empty else ByteVector(str getBytes charset)
  def encodeStringStrict(str: String, charset: Charset)   = ByteVector.encodeString(str)(charset)

  ///////////////// QUERY ///////////////////

  def size(value: ByteVector)                                           = value.size
  def byteAt(value: ByteVector, ix: Long)                               = value(ix)
  def indexOfSlice(value: ByteVector, slice: ByteVector, startIx: Long) = value.indexOfSlice(slice, startIx)

  ///////////////// TRANSFORMATION TO ByteVector ///////////////////

  def update(value: ByteVector, ix: Long, byte: Byte) = value.update(ix, byte)
  def concat(value: ByteVector, other: ByteVector)    = value ++ other
  def concat(value: ByteVector, byte: Byte)           = value :+ byte
  def concat(byte: Byte, value: ByteVector)           = byte +: value
  def drop(value: ByteVector, n: Long)                = value.drop(n)
  def take(value: ByteVector, n: Long)                = value.take(n)
  def map(value: ByteVector, f: Byte ⇒ Byte)          = value.map(f)
  def reverse(value: ByteVector)                      = value.reverse
  def compact(value: ByteVector)                      = value.compact

  ///////////////// TRANSFORMATION TO OTHER TYPES ///////////////////

  def toArray(value: ByteVector)                                   = value.toArray
  def copyToArray(value: ByteVector, xs: Array[Byte], offset: Int) = value.copyToArray(xs, offset)
  def copyToArray(value: ByteVector, sourceOffset: Long, xs: Array[Byte], destOffset: Int, len: Int) =
    value.copyToArray(xs, destOffset, sourceOffset, len)
  def copyToBuffer(value: ByteVector, buffer: ByteBuffer): Int = value.copyToBuffer(buffer)
  def copyToOutputStream(value: ByteVector, s: OutputStream)   = value.copyToStream(s)
  def toByteBuffer(value: ByteVector)                          = value.toByteBuffer
  def toIndexedSeq(value: ByteVector): IndexedSeq[Byte]        = value.toIndexedSeq
  def toSeq(value: ByteVector): Seq[Byte]                      = value.toSeq
  def decodeString(value: ByteVector, charset: Charset): Either[CharacterCodingException, String] =
    value.decodeString(charset)

  ///////////////// ITERATION ///////////////////

  def foldLeft[A](value: ByteVector, z: A, f: (A, Byte) ⇒ A)  = value.foldLeft(z)(f)
  def foldRight[A](value: ByteVector, z: A, f: (Byte, A) ⇒ A) = value.foldRight(z)(f)
  def foreach(value: ByteVector, f: Byte ⇒ Unit)              = value.foreach(f)

} 
Example 163
Source File: DataTypeCodec.scala    From polynote   with Apache License 2.0 5 votes vote down vote up
package polynote.kernel

import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets

import polynote.runtime._
import scodec.Codec
import scodec.bits.ByteVector
import scodec.codecs._
import scodec.codecs.implicits._
import shapeless.cachedImplicit

object DataTypeCodec {

  implicit val byteDiscriminator: Discriminator[DataType, ByteType.type, Byte] = Discriminator(0)
  implicit val boolDiscriminator: Discriminator[DataType, BoolType.type, Byte] = Discriminator(1)
  implicit val shortDiscriminator: Discriminator[DataType, ShortType.type, Byte] = Discriminator(2)
  implicit val intDiscriminator: Discriminator[DataType, IntType.type, Byte] = Discriminator(3)
  implicit val longDiscriminator: Discriminator[DataType, LongType.type, Byte] = Discriminator(4)
  implicit val floatDiscriminator: Discriminator[DataType, FloatType.type, Byte] = Discriminator(5)
  implicit val doubleDiscriminator: Discriminator[DataType, DoubleType.type, Byte] = Discriminator(6)
  implicit val binaryDiscriminator: Discriminator[DataType, BinaryType.type, Byte] = Discriminator(7)
  implicit val stringDiscriminator: Discriminator[DataType, StringType.type, Byte] = Discriminator(8)
  implicit val structDiscriminator: Discriminator[DataType, StructType, Byte] = Discriminator(9)
  implicit val optionalDiscriminator: Discriminator[DataType, OptionalType, Byte] = Discriminator(10)
  implicit val arrayDiscriminator: Discriminator[DataType, ArrayType, Byte] = Discriminator(11)
  //implicit val dateDiscriminator: Discriminator[DataType, DateType.type, Byte] = Discriminator(12)
  //implicit val timestampDiscriminator: Discriminator[DataType, TimestampType.type, Byte] = Discriminator(13)
  implicit val typeDiscriminator: Discriminator[DataType, TypeType.type, Byte] = Discriminator(14)
  implicit val mapDiscriminator: Discriminator[DataType, MapType, Byte] = Discriminator(15)

  implicit val dataTypeDiscriminated: Discriminated[DataType, Byte] = Discriminated(byte)

  implicit val dataTypeCodec: Codec[DataType] = cachedImplicit
}

object ValueReprCodec {
  import DataTypeCodec.dataTypeCodec

  implicit val valueReprDiscriminated: Discriminated[ValueRepr, Byte] = Discriminated(byte)
  implicit val stringRepr: Discriminator[ValueRepr, StringRepr, Byte] = Discriminator(0)
  implicit val mimeRepr: Discriminator[ValueRepr, MIMERepr, Byte] =  Discriminator(1)
  implicit val dataRepr: Discriminator[ValueRepr, DataRepr, Byte] = Discriminator(2)
  implicit val lazyDataRepr: Discriminator[ValueRepr, LazyDataRepr, Byte] = Discriminator(3)
  implicit val updatingDataRepr: Discriminator[ValueRepr, UpdatingDataRepr, Byte] = Discriminator(4)
  implicit val streamingDataRepr: Discriminator[ValueRepr, StreamingDataRepr, Byte] = Discriminator(5)

  implicit val streamingDataReprCodec: Codec[StreamingDataRepr] = cachedImplicit
  implicit val byteBufferCodec: Codec[ByteBuffer] = variableSizeBytes(int32, bytes).xmap(_.toByteBuffer, ByteVector.apply)

  implicit val codec: Codec[ValueRepr] = cachedImplicit
}

object TableOpCodec {
  implicit val tableOpDiscriminated: Discriminated[TableOp, Byte] = Discriminated(byte)
  implicit val groupAgg: Discriminator[TableOp, GroupAgg, Byte] = Discriminator(0)
  implicit val quantileBin: Discriminator[TableOp, QuantileBin, Byte] = Discriminator(1)
  implicit val select: Discriminator[TableOp, Select, Byte] = Discriminator(2)

  implicit val tableOpCodec: Codec[TableOp] = cachedImplicit
} 
Example 164
Source File: ResultOutputStream.scala    From polynote   with Apache License 2.0 5 votes vote down vote up
package polynote.kernel.util

import java.io.OutputStream
import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets
import java.util.concurrent.atomic.AtomicBoolean

import polynote.kernel.{Output, Result}

class ResultOutputStream(publishSync: Result => Unit, bufSize: Int = 1024) extends OutputStream {
  private val buf: ByteBuffer = ByteBuffer.allocate(bufSize)
  private val closed = new AtomicBoolean(false)

  def write(b: Int): Unit = buf.synchronized {
    if (!buf.hasRemaining) {
      flush()
    }

    buf.put(b.toByte)
  }

  override def flush(): Unit = {
    super.flush()
    buf.synchronized {
      val len = buf.position()
      if (len > 0) {
        val b = ByteBuffer.allocate(buf.position())
        val arr = new Array[Byte](buf.position())
        buf.rewind()
        buf.get(arr)
        publishSync(Output("text/plain; rel=stdout", new String(arr, StandardCharsets.UTF_8)))
        buf.rewind()
      }
    }
  }

  override def close(): Unit = buf.synchronized {
    if (!closed.getAndSet(true)) {
      flush()
      super.close()
    }
  }

} 
Example 165
Source File: WAL.scala    From polynote   with Apache License 2.0 5 votes vote down vote up
package polynote.server
package repository.fs

import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets
import java.time.Instant

import polynote.messages.{Message, Notebook}
import scodec.bits.{BitVector, ByteVector}
import scodec.{Attempt, Codec, codecs}
import scodec.stream.decode
import scodec.stream.decode.StreamDecoder
import zio.{RIO, Task, ZIO}
import zio.blocking.Blocking
import zio.clock.{Clock, currentDateTime}

import scala.util.Try

object WAL {

  val WALMagicNumber: Array[Byte] = "PNWAL".getBytes(StandardCharsets.US_ASCII)
  val WALVersion: Short = 1

  // Timestamp for each update is stored in 32 bits unsigned, epoch UTC seconds.
  // So we'll have to change the format by February of 2106. Apologies to my great-great-great grandchildren.
  private val instantCodec = codecs.uint32.exmap[Instant](
    epochSeconds => Attempt.fromTry(Try(Instant.ofEpochSecond(epochSeconds))),
    instant      => Attempt.successful(instant.getEpochSecond)
  )

  def encodeTimestamp(instant: Instant): Task[BitVector] =
    ZIO.fromEither(instantCodec.encode(instant).toEither)
      .mapError(err => new RuntimeException(err.message))

  val messageCodec: Codec[(Instant, Message)] = codecs.variableSizeBytes(codecs.int32, instantCodec ~ Message.codec)

  val decoder: StreamDecoder[(Instant, Message)] = {
    val readMagic = decode.once(codecs.constant(ByteVector(WALMagicNumber)))
    val readVersion = decode.once(codecs.int16)
    def readMessages(version: Int): StreamDecoder[(Instant, Message)] = version match {
      case 1 => decode.many(messageCodec)
      case v => decode.raiseError(new Exception(s"Unknown WAL version $v"))
    }

    for {
      _       <- readMagic
      ver     <- readVersion
      message <- readMessages(ver)
    } yield message
  }

  trait WALWriter {
    protected def append(bytes: Array[Byte]): RIO[Blocking, Unit] = append(ByteBuffer.wrap(bytes))
    protected def append(bytes: BitVector): RIO[Blocking, Unit] = append(bytes.toByteBuffer)
    protected def append(bytes: ByteBuffer): RIO[Blocking, Unit]

    def writeHeader(notebook: Notebook): RIO[Blocking with Clock, Unit] =
      append(WALMagicNumber) *>
        append(BitVector.fromShort(WALVersion)) *>
        appendMessage(notebook.withoutResults)

    def appendMessage(message: Message): RIO[Blocking with Clock, Unit] = for {
      ts    <- currentDateTime.map(_.toInstant)
      bytes <- ZIO.fromEither(messageCodec.encode((ts, message)).toEither).mapError(err => new RuntimeException(err.message))
      _     <- append(bytes)
    } yield ()

    def sync(): RIO[Blocking, Unit]

    def close(): RIO[Blocking, Unit]
  }

  object WALWriter {
    object NoWAL extends WALWriter {
      override protected def append(bytes: ByteBuffer): RIO[Blocking, Unit] = ZIO.unit
      override def sync(): RIO[Blocking, Unit] = ZIO.unit
      override def close(): RIO[Blocking, Unit] = ZIO.unit
    }
  }
} 
Example 166
Source File: CollectionReprsSpec.scala    From polynote   with Apache License 2.0 5 votes vote down vote up
package polynote.runtime.test

import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets

import org.scalatest.{FreeSpec, Matchers}
import polynote.runtime.{DataEncoder, GroupAgg, ReprsOf, StreamingDataRepr}

class  CollectionReprsSpec extends FreeSpec with Matchers {

  "Streaming repr of structs" - {
    case class Example(label: String, i: Int, d: Double)

    "Aggregates correctly" - {
      "mean" in {
        val l = List(Example("a", 10, 10.0), Example("b", 11, 11.0), Example("c", 12, 12.0), Example("a", 12, 12.0))
        val de = implicitly[DataEncoder.StructDataEncoder[Example]]
        val h = ReprsOf.StructSeqStreamHandle[Example, Example](0, l, l => l, de)
        val Right(h1) = h.modify(List(GroupAgg(List("label"), List("i" -> "mean", "d" -> "mean")))).right.map(_.apply(1))

        def decode(buf: ByteBuffer) = {
          buf.rewind()
          val labelLength = buf.getInt()
          val labelArr = new Array[Byte](labelLength)
          buf.get(labelArr)
          val label = new String(labelArr, StandardCharsets.UTF_8)
          val avgI = buf.getDouble()
          val avgD = buf.getDouble()

          (label, avgI, avgD)
        }

        h1.iterator.map(decode).toList should contain theSameElementsAs List(
          ("a", 11.0, 11.0),
          ("b", 11.0, 11.0),
          ("c", 12.0, 12.0)
        )
      }

    }
  }

} 
Example 167
Source File: XORShiftRandom.scala    From Mastering-Spark-for-Data-Science   with MIT License 5 votes vote down vote up
package io.gzet.story.linalg

import java.nio.ByteBuffer
import java.util.{Random => JavaRandom}

import scala.util.Random
import scala.util.hashing.MurmurHash3


  private[this] def hashSeed(seed: Long): Long = {
    val bytes = ByteBuffer.allocate(java.lang.Long.SIZE).putLong(seed).array()
    MurmurHash3.bytesHash(bytes)
  }

  // we need to just override next - this will be called by nextInt, nextDouble,
  // nextGaussian, nextLong, etc.
  override protected def next(bits: Int): Int = {
    var nextSeed = seed ^ (seed << 21)
    nextSeed ^= (nextSeed >>> 35)
    nextSeed ^= (nextSeed << 4)
    seed = nextSeed
    (nextSeed & ((1L << bits) - 1)).asInstanceOf[Int]
  }
}

object XORShiftRandom {
  val random = new Random()
} 
Example 168
Source File: IndexCodec.scala    From unicorn   with Apache License 2.0 5 votes vote down vote up
package unicorn.index

import java.nio.ByteBuffer

import unicorn.bigtable.{Cell, Column}
import unicorn.index.IndexSortOrder._
import unicorn.util._


  def resetBuffer(tenant: Option[Array[Byte]]): Unit = {
    buffer.clear
    buffer.putShort(index.id.toShort)
    tenant match {
      case None => buffer.put(0.toByte)
      case Some(tenant) => buffer.put(tenant.length.toByte).put(tenant)
    }
  }
}

object IndexCodec {
  def apply(index: Index): IndexCodec = {
    index.indexType match {
      case IndexType.Hashed => new HashIndexCodec(index)
      case IndexType.Text => new TextIndexCodec(index)
      case _ => if (index.columns.size == 1) new SingleColumnIndexCodec(index) else new CompositeIndexCodec(index)
    }
  }
} 
Example 169
Source File: HashIndexCodec.scala    From unicorn   with Apache License 2.0 5 votes vote down vote up
package unicorn.index

import java.nio.ByteBuffer

import unicorn.bigtable.Cell
import unicorn.index.IndexSortOrder._
import unicorn.util._


class HashIndexCodec(val index: Index) extends IndexCodec {
  require(index.indexType == IndexType.Hashed)

  override def apply(tenant: Option[Array[Byte]], row: ByteArray, columns: ColumnMap): Seq[Cell] = {
    val hasUndefinedColumn = index.columns.exists { indexColumn =>
      columns.get(index.family).map(_.get(indexColumn.qualifier)).getOrElse(None) match {
        case Some(_) => false
        case None => true
      }
    }

    if (hasUndefinedColumn) return Seq.empty

    val hasZeroTimestamp = index.columns.exists { indexColumn =>
      columns.get(index.family).map(_.get(indexColumn.qualifier)).getOrElse(None) match {
        case Some(c) => c.timestamp == 0L
        case None => false
      }
    }

    val timestamp = if (hasZeroTimestamp) 0L else index.columns.foldLeft(0L) { (b, indexColumn) =>
      val ts = columns(index.family)(indexColumn.qualifier).timestamp
      Math.max(b, ts)
    }

    resetBuffer(tenant)
    index.columns.foreach { indexColumn =>
      val column = columns(index.family)(indexColumn.qualifier).value

      indexColumn.order match {
        case Ascending => buffer.put(column)
        case Descending => buffer.put(~column)
      }
    }

    val key = md5(buffer)

    val (qualifier: ByteArray, indexValue: ByteArray) = index.indexType match {
      case IndexType.Unique => (UniqueIndexColumnQualifier, row)
      case _ => (row, IndexDummyValue)
    }

    Seq(Cell(key, IndexColumnFamily, qualifier, indexValue, timestamp))
  }
} 
Example 170
Source File: DocumentSerializer.scala    From unicorn   with Apache License 2.0 5 votes vote down vote up
package unicorn.unibase

import java.nio.ByteBuffer

import unicorn.bigtable.{Column, ColumnFamily}
import unicorn.json._


  def deserialize(data: Seq[ColumnFamily]): Option[JsObject] = {
    val objects = data.map { case ColumnFamily(family, columns) =>
      val map = columns.map { case Column(qualifier, value, _) =>
        (new String(qualifier, valueSerializer.charset), value.bytes)
      }.toMap
      val json = valueSerializer.deserialize(map)
      json.asInstanceOf[JsObject]
    }

    if (objects.size == 0)
      None
    else if (objects.size == 1)
      Some(objects(0))
    else {
      val fold = objects.foldLeft(JsObject()) { (doc, family) =>
        doc.fields ++= family.fields
        doc
      }
      Some(fold)
    }
  }
} 
Example 171
Source File: ObjectId.scala    From unicorn   with Apache License 2.0 5 votes vote down vote up
package unicorn.oid

import java.util.{Arrays, Date, UUID}
import java.nio.ByteBuffer
import unicorn.util._


  def string: String = {
    new String(id, utf8)
  }
}

object ObjectId {
  def apply(id: Array[Byte]) = new ObjectId(id)
  def apply(id: String) = new ObjectId(id.getBytes("UTF-8"))

  def apply(id: Int) = {
    val array = Array[Byte](4)
    val buffer = ByteBuffer.wrap(array)
    buffer.putInt(id)
    new ObjectId(array)
  }

  def apply(id: Long) = {
    val array = Array[Byte](8)
    val buffer = ByteBuffer.wrap(array)
    buffer.putLong(id)
    new ObjectId(array)
  }

  def apply(id: Date) = {
    val array = Array[Byte](8)
    val buffer = ByteBuffer.wrap(array)
    buffer.putLong(id.getTime)
    new ObjectId(array)
  }

  def apply(id: UUID) = {
    val array = Array[Byte](16)
    val buffer = ByteBuffer.wrap(array)
    buffer.putLong(id.getMostSignificantBits)
    buffer.putLong(id.getLeastSignificantBits)
    new ObjectId(array)
  }
} 
Example 172
Source File: FieldDateTime.scala    From spark-gdb   with Apache License 2.0 5 votes vote down vote up
package com.esri.gdb

import java.nio.ByteBuffer
import java.sql.Timestamp

import org.apache.spark.sql.types.{Metadata, TimestampType}


class FieldDateTime(name: String, nullValueAllowed: Boolean, metadata:Metadata)
  extends Field(name, TimestampType, nullValueAllowed, metadata) {

  override def readValue(byteBuffer: ByteBuffer, oid: Int) = {
    val numDays = byteBuffer.getDouble
    // convert days since 12/30/1899 to 1/1/1970
    val unixDays = numDays - 25569
    val millis = (unixDays * 1000 * 60 * 60 * 24).ceil.toLong
    new Timestamp(millis)
  }
} 
Example 173
Source File: FieldPointMType.scala    From spark-gdb   with Apache License 2.0 5 votes vote down vote up
package com.esri.gdb

import java.nio.ByteBuffer

import com.esri.udt.{PointMType, PointMUDT}
import org.apache.spark.sql.types.Metadata


object FieldPointMType extends Serializable {
  def apply(name: String,
            nullValueAllowed: Boolean,
            xOrig: Double,
            yOrig: Double,
            mOrig: Double,
            xyScale: Double,
            mScale: Double,
            metadata: Metadata
           ) = {
    new FieldPointMType(name, nullValueAllowed, xOrig, yOrig, mOrig, xyScale, mScale, metadata)
  }
}

class FieldPointMType(name: String,
                      nullValueAllowed: Boolean,
                      xOrig: Double,
                      yOrig: Double,
                      mOrig: Double,
                      xyScale: Double,
                      mScale: Double,
                      metadata: Metadata)
  extends FieldBytes(name, new PointMUDT(), nullValueAllowed, metadata) {

  override def readValue(byteBuffer: ByteBuffer, oid: Int) = {
    val blob = getByteBuffer(byteBuffer)

    val geomType = blob.getVarUInt()

    val vx = blob.getVarUInt
    val vy = blob.getVarUInt
    val vm = blob.getVarUInt

    val x = (vx - 1.0) / xyScale + xOrig
    val y = (vy - 1.0) / xyScale + yOrig
    val m = (vm - 1.0) / mScale + mOrig

    new PointMType(x, y, m)
  }
} 
Example 174
Source File: GDBIndex.scala    From spark-gdb   with Apache License 2.0 5 votes vote down vote up
package com.esri.gdb

import java.io.{DataInput, File}
import java.nio.{ByteBuffer, ByteOrder}

import org.apache.hadoop.conf.Configuration
import org.apache.hadoop.fs.{FSDataInputStream, Path}
import org.apache.spark.Logging

object GDBIndex {
  def apply(path: String, name: String, conf: Configuration = new Configuration()) = {
    val filename = StringBuilder.newBuilder.append(path).append(File.separator).append(name).append(".gdbtablx").toString()
    val hdfsPath = new Path(filename)
    val dataInput = hdfsPath.getFileSystem(conf).open(hdfsPath)

    val bytes = new Array[Byte](16)
    dataInput.readFully(bytes)
    val byteBuffer = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN)

    val signature = byteBuffer.getInt
    val n1024Blocks = byteBuffer.getInt
    val numRows = byteBuffer.getInt
    val indexSize = byteBuffer.getInt

    new GDBIndex(dataInput, numRows, indexSize)
  }
}

private[gdb] class GDBIndex(dataInput: FSDataInputStream,
                            val numRows: Int,
                            indexSize: Int
                           ) extends Logging with AutoCloseable with Serializable {

  def readSeekForRowNum(rowNum: Int) = {
    val bytes = new Array[Byte](indexSize)
    dataInput.seek(16 + rowNum * indexSize)
    dataInput.readFully(bytes)
    ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN).getInt
  }

  def iterator(startAtRow: Int = 0, numRowsToRead: Int = -1) = {
    dataInput.seek(16 + startAtRow * indexSize)
    val maxRows = if (numRowsToRead == -1) numRows else numRowsToRead
    // log.info(s"iterator::startAtRow=$startAtRow maxRows=$maxRows")
    new GDBIndexIterator(dataInput, startAtRow, maxRows, indexSize).withFilter(_.isSeekable)
  }

  def close() {
    dataInput.close()
  }
}

private[gdb] class GDBIndexIterator(dataInput: DataInput,
                                    startID: Int,
                                    maxRows: Int,
                                    indexSize: Int
                                   ) extends Iterator[IndexInfo] with Logging with Serializable {

  private val indexInfo = IndexInfo(0, 0)
  private val bytes = new Array[Byte](indexSize)
  private val byteBuffer = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN)

  private var objectID = startID
  private var nextRow = 0

  def hasNext() = nextRow < maxRows

  def next() = {
    // log.info(s"next::nextRow=$nextRow maxRows=$maxRows")
    nextRow += 1

    objectID += 1
    indexInfo.objectID = objectID

    byteBuffer.clear
    dataInput.readFully(bytes)
    indexInfo.seek = byteBuffer.getInt

    indexInfo
  }
} 
Example 175
Source File: FieldPoly.scala    From spark-gdb   with Apache License 2.0 5 votes vote down vote up
package com.esri.gdb

import java.nio.ByteBuffer

import com.esri.core.geometry.MultiPath
import org.apache.spark.sql.types.{DataType, Metadata}

@deprecated("not used", "0.4")
abstract class FieldPoly(name: String,
                         dataType: DataType,
                         nullValueAllowed: Boolean,
                         xOrig: Double,
                         yOrig: Double,
                         xyScale: Double,
                         metadata: Metadata)
  extends FieldBytes(name, dataType, nullValueAllowed, metadata) {

  protected var dx = 0L
  protected var dy = 0L

  def addPath(byteBuffer: ByteBuffer, numCoordinates: Int, path: MultiPath) = {
    0 until numCoordinates foreach (n => {
      dx += byteBuffer.getVarInt
      dy += byteBuffer.getVarInt
      val x = dx / xyScale + xOrig
      val y = dy / xyScale + yOrig
      n match {
        case 0 => path.startPath(x, y)
        case _ => path.lineTo(x, y)
      }
    })
    path
  }
} 
Example 176
Source File: FieldBytes.scala    From spark-gdb   with Apache License 2.0 5 votes vote down vote up
package com.esri.gdb

import java.nio.ByteBuffer

import org.apache.spark.sql.types.{DataType, Metadata}


abstract class FieldBytes(name: String,
                          dataType: DataType,
                          nullValueAllowed: Boolean,
                          metadata: Metadata = Metadata.empty
                         )
  extends Field(name, dataType, nullValueAllowed, metadata) {

  protected var m_bytes = new Array[Byte](1024)

  def getByteBuffer(byteBuffer: ByteBuffer) = {
    val numBytes = fillVarBytes(byteBuffer)
    ByteBuffer.wrap(m_bytes, 0, numBytes)
  }

  def fillVarBytes(byteBuffer: ByteBuffer) = {
    val numBytes = byteBuffer.getVarUInt.toInt
    if (numBytes > m_bytes.length) {
      m_bytes = new Array[Byte](numBytes)
    }
    0 until numBytes foreach {
      m_bytes(_) = byteBuffer.get
    }
    numBytes
  }
} 
Example 177
Source File: FieldPolygon.scala    From spark-gdb   with Apache License 2.0 5 votes vote down vote up
package com.esri.gdb

import java.nio.ByteBuffer

import com.esri.core.geometry.Polygon
import com.esri.udt.PolygonUDT
import org.apache.spark.sql.types.{DataType, Metadata}

@deprecated("not used", "0.4")
object FieldPolygon {
  def apply(name: String,
            nullValueAllowed: Boolean,
            xOrig: Double,
            yOrig: Double,
            xyScale: Double,
            metadata: Metadata) = {
    new FieldPolygonEsri(name, nullValueAllowed, xOrig, yOrig, xyScale, metadata)
  }
}

@deprecated("not used", "0.4")
abstract class FieldPolygon(name: String,
                            dataType: DataType,
                            nullValueAllowed: Boolean,
                            xOrig: Double,
                            yOrig: Double,
                            xyScale: Double,
                            metadata: Metadata
                           )
  extends FieldPoly(name, dataType, nullValueAllowed, xOrig, yOrig, xyScale, metadata) {

  override def readValue(byteBuffer: ByteBuffer, oid: Int) = {
    val polygon = new Polygon()

    val blob = getByteBuffer(byteBuffer)

    val geomType = blob.getVarUInt

    val numPoints = blob.getVarUInt.toInt
    val numParts = blob.getVarUInt.toInt

    val xmin = blob.getVarUInt / xyScale + xOrig
    val ymin = blob.getVarUInt / xyScale + yOrig
    val xmax = blob.getVarUInt / xyScale + xmin
    val ymax = blob.getVarUInt / xyScale + ymin

    dx = 0L
    dy = 0L

    if (numParts > 1) {
      var sum = 0
      val numCoordSeq = 1 to numParts map (part => {
        val numCoord = if (part == numParts) {
          numPoints - sum
        } else {
          blob.getVarUInt.toInt
        }
        sum += numCoord
        numCoord
      })
      // TODO - fix shells and holes based on https://github.com/rouault/dump_gdbtable/wiki/FGDB-Spec
      numCoordSeq.foreach(numCoord => addPath(blob, numCoord, polygon))
    }
    else {
      addPath(blob, numPoints, polygon)
    }
    polygon
  }
}

@deprecated("not used", "0.4")
class FieldPolygonEsri(name: String,
                       nullValueAllowed: Boolean,
                       xOrig: Double,
                       yOrig: Double,
                       xyScale: Double,
                       metadata: Metadata)
  extends FieldPolygon(name, new PolygonUDT(), nullValueAllowed, xOrig, yOrig, xyScale, metadata) 
Example 178
Source File: FieldPoly2Type.scala    From spark-gdb   with Apache License 2.0 5 votes vote down vote up
package com.esri.gdb

import java.nio.ByteBuffer

import org.apache.spark.sql.types.{DataType, Metadata}

abstract class FieldPoly2Type[T](name: String,
                                 dataType: DataType,
                                 nullValueAllowed: Boolean,
                                 xOrig: Double,
                                 yOrig: Double,
                                 xyScale: Double,
                                 metadata: Metadata)
  extends FieldBytes(name, dataType, nullValueAllowed, metadata) {

  override def readValue(byteBuffer: ByteBuffer, oid: Int) = {
    val blob = getByteBuffer(byteBuffer)

    val geomType = blob.getVarUInt
    val numPoints = blob.getVarUInt.toInt
    if (numPoints == 0)
      createPolyType(0, 0, 0, 0, Array.empty[Int], Array.empty[Double])
    else {
      val numParts = blob.getVarUInt.toInt

      val xmin = blob.getVarUInt / xyScale + xOrig
      val ymin = blob.getVarUInt / xyScale + yOrig
      val xmax = blob.getVarUInt / xyScale + xmin
      val ymax = blob.getVarUInt / xyScale + ymin

      var dx = 0L
      var dy = 0L

      val xyNum = new Array[Int](numParts)
      val xyArr = new Array[Double](numPoints * 2)

      if (numParts > 1) {
        var i = 0
        var sum = 0
        1 to numParts foreach (partIndex => {
          if (partIndex == numParts) {
            xyNum(i) = numPoints - sum
          } else {
            val numXY = blob.getVarUInt.toInt
            xyNum(i) = numXY
            sum += numXY
            i += 1
          }
        })
        i = 0
        xyNum.foreach(numXY => {
          0 until numXY foreach (n => {
            dx += blob.getVarInt
            dy += blob.getVarInt
            val x = dx / xyScale + xOrig
            val y = dy / xyScale + yOrig
            xyArr(i) = x
            i += 1
            xyArr(i) = y
            i += 1
          })
        })
      }
      else {
        xyNum(0) = numPoints
        var i = 0
        0 until numPoints foreach (n => {
          dx += blob.getVarInt
          dy += blob.getVarInt
          val x = dx / xyScale + xOrig
          val y = dy / xyScale + yOrig
          xyArr(i) = x
          i += 1
          xyArr(i) = y
          i += 1
        })
      }
      createPolyType(xmin, ymin, xmax, ymax, xyNum, xyArr)
    }
  }

  def createPolyType(xmin: Double, ymin: Double, xmax: Double, ymax: Double, xyNum: Array[Int], xyArr: Array[Double]): T
} 
Example 179
Source File: FieldPolyline.scala    From spark-gdb   with Apache License 2.0 5 votes vote down vote up
package com.esri.gdb

import java.nio.ByteBuffer

import com.esri.core.geometry.Polyline
import com.esri.udt.PolylineUDT
import org.apache.spark.sql.types.Metadata

@deprecated("not used", "0.4")
object FieldPolyline extends Serializable {
  def apply(name: String,
            nullValueAllowed: Boolean,
            xOrig: Double,
            yOrig: Double,
            xyScale: Double,
            metadata: Metadata) = {
    new FieldPolyline(name, nullValueAllowed, xOrig, yOrig, xyScale, metadata)
  }
}

@deprecated("not used", "0.4")
class FieldPolyline(name: String,
                    nullValueAllowed: Boolean,
                    xOrig: Double,
                    yOrig: Double,
                    xyScale: Double,
                    metadata: Metadata
                   )
  extends FieldPoly(name, new PolylineUDT(), nullValueAllowed, xOrig, yOrig, xyScale, metadata) {

  override def readValue(byteBuffer: ByteBuffer, oid: Int) = {
    val polyline = new Polyline()

    val blob = getByteBuffer(byteBuffer)
    val geomType = blob.getVarUInt

    val numPoints = blob.getVarUInt.toInt
    val numParts = blob.getVarUInt.toInt

    val xmin = blob.getVarUInt / xyScale + xOrig
    val ymin = blob.getVarUInt / xyScale + yOrig
    val xmax = blob.getVarUInt / xyScale + xmin
    val ymax = blob.getVarUInt / xyScale + ymin

    dx = 0L
    dy = 0L

    if (numParts > 1) {
      var sum = 0
      val numCoordSeq = 1 to numParts map (part => {
        val numCoord = if (part == numParts) {
          numPoints - sum
        } else {
          blob.getVarUInt.toInt
        }
        sum += numCoord
        numCoord
      })
      numCoordSeq.foreach(numCoord => addPath(blob, numCoord, polyline))
    }
    else {
      addPath(blob, numPoints, polyline)
    }
    polyline
  }
} 
Example 180
Source File: FieldPointZMType.scala    From spark-gdb   with Apache License 2.0 5 votes vote down vote up
package com.esri.gdb

import java.nio.ByteBuffer

import com.esri.udt.{PointZMType, PointZMUDT}
import org.apache.spark.sql.types.Metadata


object FieldPointZMType extends Serializable {
  def apply(name: String,
            nullValueAllowed: Boolean,
            xOrig: Double,
            yOrig: Double,
            zOrig: Double,
            mOrig: Double,
            xyScale: Double,
            zScale: Double,
            mScale: Double,
            metadata: Metadata
           ) = {
    new FieldPointZMType(name, nullValueAllowed, xOrig, yOrig, zOrig, mOrig, xyScale, zScale, mScale, metadata)
  }
}

class FieldPointZMType(name: String,
                       nullValueAllowed: Boolean,
                       xOrig: Double,
                       yOrig: Double,
                       zOrig: Double,
                       mOrig: Double,
                       xyScale: Double,
                       zScale: Double,
                       mScale: Double,
                       metadata: Metadata)
  extends FieldBytes(name, new PointZMUDT(), nullValueAllowed, metadata) {

  override def readValue(byteBuffer: ByteBuffer, oid: Int) = {
    val blob = getByteBuffer(byteBuffer)

    val geomType = blob.getVarUInt

    val vx = blob.getVarUInt()
    val vy = blob.getVarUInt()
    val x = (vx - 1.0) / xyScale + xOrig
    val y = (vy - 1.0) / xyScale + yOrig

    geomType match {
      // Point
      case 1 => new PointZMType(x, y)
      // PointZ
      case 9 =>
        val vz = blob.getVarUInt
        val z = (vz - 1.0) / zScale + zOrig
        new PointZMType(x, y, z)
      // PointM
      case 21 =>
        val vm = blob.getVarUInt
        val m = (vm - 1.0) / mScale + mOrig
        new PointZMType(x, y, 0.0, m)
      // PointZM
      case _ =>
        val vz = blob.getVarUInt
        val vm = blob.getVarUInt
        val z = (vz - 1.0) / zScale + zOrig
        val m = (vm - 1.0) / mScale + mOrig
        new PointZMType(x, y, z, m)
    }
  }
} 
Example 181
Source File: DataBuffer.scala    From spark-gdb   with Apache License 2.0 5 votes vote down vote up
package com.esri.gdb

import java.nio.{ByteBuffer, ByteOrder}

import org.apache.hadoop.fs.FSDataInputStream


class DataBuffer(dataInput: FSDataInputStream) extends Serializable {

  private var bytes = new Array[Byte](1024)
  private var byteBuffer = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN)

  def readBytes(length: Int) = {
    if (length > bytes.length) {
      bytes = new Array[Byte](length)
      byteBuffer = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN)
    }
    else {
      byteBuffer.clear
    }
    dataInput.readFully(bytes, 0, length)
    byteBuffer
  }

  def seek(position: Long) = {
    dataInput.seek(position)
    this
  }

  def close() {
    dataInput.close()
  }
}

object DataBuffer {
  def apply(dataInput: FSDataInputStream) = {
    new DataBuffer(dataInput)
  }
} 
Example 182
Source File: FieldPointZType.scala    From spark-gdb   with Apache License 2.0 5 votes vote down vote up
package com.esri.gdb

import java.nio.ByteBuffer

import com.esri.udt.{PointZType, PointZUDT}
import org.apache.spark.sql.types.Metadata


object FieldPointZType extends Serializable {
  def apply(name: String,
            nullValueAllowed: Boolean,
            xOrig: Double,
            yOrig: Double,
            zOrig: Double,
            xyScale: Double,
            zScale: Double,
            metadata: Metadata
           ) = {
    new FieldPointZType(name, nullValueAllowed, xOrig, yOrig, zOrig, xyScale, zScale, metadata)
  }
}

class FieldPointZType(name: String,
                      nullValueAllowed: Boolean,
                      xOrig: Double,
                      yOrig: Double,
                      zOrig: Double,
                      xyScale: Double,
                      zScale: Double,
                      metadata: Metadata)
  extends FieldBytes(name, new PointZUDT(), nullValueAllowed, metadata) {

  override def readValue(byteBuffer: ByteBuffer, oid: Int) = {
    val blob = getByteBuffer(byteBuffer)

    val geomType = blob.getVarUInt

    val vx = blob.getVarUInt
    val vy = blob.getVarUInt
    val vz = blob.getVarUInt

    val x = (vx - 1.0) / xyScale + xOrig
    val y = (vy - 1.0) / xyScale + yOrig
    val z = (vz - 1.0) / zScale + zOrig

    new PointZType(x, y, z)
  }
} 
Example 183
Source File: FieldPoly3Type.scala    From spark-gdb   with Apache License 2.0 5 votes vote down vote up
package com.esri.gdb

import java.nio.ByteBuffer

import org.apache.spark.sql.types.{DataType, Metadata}

abstract class FieldPoly3Type[T](name: String,
                                 dataType: DataType,
                                 nullValueAllowed: Boolean,
                                 xOrig: Double,
                                 yOrig: Double,
                                 nOrig: Double,
                                 xyScale: Double,
                                 nScale: Double,
                                 metadata: Metadata)
  extends FieldBytes(name, dataType, nullValueAllowed, metadata) {

  override def readValue(byteBuffer: ByteBuffer, oid: Int) = {
    val blob = getByteBuffer(byteBuffer)

    val geomType = blob.getVarUInt
    val numPoints = blob.getVarUInt.toInt
    // TODO - Handle zero num points in other geom type.
    if (numPoints == 0) {
      createPolyMType(0, 0, 0, 0, Array.empty[Int], Array.empty[Double])
    }
    else {
      val numParts = blob.getVarUInt.toInt

      val xmin = blob.getVarUInt / xyScale + xOrig
      val ymin = blob.getVarUInt / xyScale + yOrig
      val xmax = blob.getVarUInt / xyScale + xmin
      val ymax = blob.getVarUInt / xyScale + ymin

      var dx = 0L
      var dy = 0L

      val xyNum = new Array[Int](numParts)
      val xyArr = new Array[Double](numPoints * 3)

      var i = 0
      if (numParts > 1) {
        var sum = 0
        1 to numParts foreach (partIndex => {
          if (partIndex == numParts) {
            xyNum(i) = numPoints - sum
          } else {
            val numXY = blob.getVarUInt.toInt
            xyNum(i) = numXY
            sum += numXY
            i += 1
          }
        })
        i = 0
        xyNum.foreach(numXY => {
          0 until numXY foreach (_ => {
            dx += blob.getVarInt
            dy += blob.getVarInt
            val x = dx / xyScale + xOrig
            val y = dy / xyScale + yOrig
            xyArr(i) = x
            i += 1
            xyArr(i) = y
            i += 2
          })
        })
      }
      else {
        xyNum(0) = numPoints
        0 until numPoints foreach (_ => {
          dx += blob.getVarInt
          dy += blob.getVarInt
          xyArr(i) = dx / xyScale + xOrig
          i += 1
          xyArr(i) = dy / xyScale + yOrig
          i += 2
        })
      }
      i = 2
      var dn = 0L
      0 until numPoints foreach (_ => {
        dn += blob.getVarInt
        xyArr(i) = dn / nScale + nOrig
        i += 3
      })
      createPolyMType(xmin, ymin, xmax, ymax, xyNum, xyArr)
    }
  }

  def createPolyMType(xmin: Double, ymin: Double, xmax: Double, ymax: Double, xyNum: Array[Int], xyArr: Array[Double]): T
} 
Example 184
Source File: FieldPointType.scala    From spark-gdb   with Apache License 2.0 5 votes vote down vote up
package com.esri.gdb

import java.nio.ByteBuffer

import com.esri.udt.{PointType, PointUDT}
import org.apache.spark.sql.types.Metadata

object FieldPointType extends Serializable {
  def apply(name: String,
            nullValueAllowed: Boolean,
            xOrig: Double,
            yOrig: Double,
            xyScale: Double,
            metadata: Metadata) = {
    new FieldPointType(name, nullValueAllowed, xOrig, yOrig, xyScale, metadata)
  }
}

class FieldPointType(name: String,
                     nullValueAllowed: Boolean,
                     xOrig: Double,
                     yOrig: Double,
                     xyScale: Double,
                     metadata: Metadata)
  extends FieldBytes(name, new PointUDT(), nullValueAllowed, metadata) {

  override def readValue(byteBuffer: ByteBuffer, oid: Int) = {
    val blob = getByteBuffer(byteBuffer)

    blob.getVarUInt() // geomType

    val vx = blob.getVarUInt()
    val vy = blob.getVarUInt()
    val x = (vx - 1.0) / xyScale + xOrig
    val y = (vy - 1.0) / xyScale + yOrig

    new PointType(x, y)
  }
} 
Example 185
Source File: package.scala    From spark-gdb   with Apache License 2.0 5 votes vote down vote up
package com.esri

import java.nio.ByteBuffer

import org.apache.spark.SparkContext
import org.apache.spark.sql.SQLContext


package object gdb {

  implicit class ByteBufferImplicits(byteBuffer: ByteBuffer) {

    implicit def getVarUInt() = {
      var shift = 7
      var b: Long = byteBuffer.get
      var ret = b & 0x7FL
      var old = ret
      while ((b & 0x80L) != 0) {
        b = byteBuffer.get
        ret = ((b & 0x7FL) << shift) | old
        old = ret
        shift += 7
      }
      ret
    }

    implicit def getVarInt() = {
      var shift = 7
      var b: Long = byteBuffer.get
      val isNeg = (b & 0x40L) != 0
      var ret = b & 0x3FL
      var old = ret
      while ((b & 0x80L) != 0) {
        b = byteBuffer.get
        ret = ((b & 0x7FL) << (shift - 1)) | old
        old = ret
        shift += 7
      }
      if (isNeg) -ret else ret
    }
  }

  implicit class SparkContextImplicits(sc: SparkContext) {
    implicit def gdbFile(path: String, name: String, numPartitions: Int = 8) = {
      GDBRDD(sc, path, name, numPartitions)
    }
  }

  implicit class SQLContextImplicits(sqlContext: SQLContext) extends Serializable {
    implicit def gdbFile(path: String, name: String, numPartitions: Int = 8) = {
      sqlContext.baseRelationToDataFrame(GDBRelation(path, name, numPartitions)(sqlContext))
    }
  }

} 
Example 186
Source File: SparkPredictionTrainer.scala    From smart-meter   with MIT License 5 votes vote down vote up
package com.logimethods.nats.connector.spark.app

import java.util.Properties;
import java.io.File
import java.io.Serializable

import org.apache.spark.SparkConf
import org.apache.spark.SparkContext
import org.apache.spark.storage.StorageLevel;
import org.apache.spark.streaming._

import io.nats.client.ConnectionFactory._
import java.nio.ByteBuffer

import org.apache.log4j.{Level, LogManager, PropertyConfigurator}

import com.logimethods.connector.nats.to_spark._
import com.logimethods.scala.connector.spark.to_nats._

import org.apache.spark.ml.classification.MultilayerPerceptronClassifier
import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator

import java.util.function._

import java.time.{LocalDateTime, ZoneOffset}
import java.time.DayOfWeek._

import org.apache.spark.ml.classification.MultilayerPerceptronClassificationModel

object SparkPredictionTrainer extends App with SparkPredictionProcessor {
  log.setLevel(Level.WARN)

  val (properties, targets, logLevel, sc, inputNatsStreaming, inputSubject, outputSubject, clusterId, outputNatsStreaming, natsUrl) = setup(args)

  val streamingDuration = scala.util.Properties.envOrElse("STREAMING_DURATION", "2000").toInt
  println("STREAMING_DURATION = " + streamingDuration)

  new Thread(new Runnable {
              def run() {
                 while( true ){
                   try {
                     val data = SparkPredictionProcessor.getData(sc, THRESHOLD)
                     val model = trainer.fit(data)
                     model.write.overwrite.save(PREDICTION_MODEL_PATH)
                     println("New model of size " + data.count() + " trained: " + model.uid)
                     Thread.sleep(streamingDuration)
                   } catch {
                     case e: Throwable => log.error(e)
                   }
                 }
              }
             }).start()
} 
Example 187
Source File: SparkProcessor.scala    From smart-meter   with MIT License 5 votes vote down vote up
package com.logimethods.nats.connector.spark.app

import java.util.Properties;
import java.io.File
import java.io.Serializable

import org.apache.spark.SparkConf
import org.apache.spark.SparkContext
import org.apache.spark.storage.StorageLevel;
import org.apache.spark.streaming._

import io.nats.client.Nats._
import io.nats.client.ConnectionFactory._
import java.nio.ByteBuffer

import org.apache.log4j.{Level, LogManager, PropertyConfigurator}

import com.logimethods.connector.nats.to_spark._
import com.logimethods.scala.connector.spark.to_nats._

import java.util.function._

import java.time.{LocalDateTime, ZoneOffset}

trait SparkProcessor {
  def setup(args: Array[String]) = {
    val inputSubject = args(0)
//    val inputNatsStreaming = inputSubject.toUpperCase.contains("STREAMING")
    val outputSubject = args(1)
//    val outputNatsStreaming = outputSubject.toUpperCase.contains("STREAMING")
    println("Will process messages from '" + inputSubject + "' to '" + outputSubject + "'")

    val logLevel = scala.util.Properties.envOrElse("LOG_LEVEL", "INFO")
    println("LOG_LEVEL = " + logLevel)

    val targets = scala.util.Properties.envOrElse("TARGETS", "ALL")
    println("TARGETS = " + targets)

    val cassandraUrl = System.getenv("CASSANDRA_URL")
    println("CASSANDRA_URL = " + cassandraUrl)

    val sparkMasterUrl = System.getenv("SPARK_MASTER_URL")
    println("SPARK_MASTER_URL = " + sparkMasterUrl)

    val sparkCoresMax = System.getenv("SPARK_CORES_MAX")
    println("SPARK_CORES_MAX = " + sparkCoresMax)

    val conf = new SparkConf()
                  .setAppName(args(2))
                  .setMaster(sparkMasterUrl)
                  .set("spark.cores.max", sparkCoresMax)
                  .set("spark.cassandra.connection.host", cassandraUrl);
    val sc = new SparkContext(conf);

//    val streamingDuration = scala.util.Properties.envOrElse("STREAMING_DURATION", "2000").toInt
//    val ssc = new StreamingContext(sc, new Duration(streamingDuration));
///    ssc.checkpoint("/spark/storage")

    val properties = new Properties();
    val natsUrl = System.getenv("NATS_URI")
    println("NATS_URI = " + natsUrl)
    properties.put("servers", natsUrl)
    properties.put(PROP_URL, natsUrl)

    val clusterId = System.getenv("NATS_CLUSTER_ID")

    val inputNatsStreaming = inputSubject.toUpperCase.contains("STREAMING")
    val outputNatsStreaming = outputSubject.toUpperCase.contains("STREAMING")

    (properties, targets, logLevel, sc, inputNatsStreaming, inputSubject, outputSubject, clusterId, outputNatsStreaming, natsUrl)
  }

  def dataDecoder: Array[Byte] => Tuple2[Long,Float] = bytes => {
        val buffer = ByteBuffer.wrap(bytes);
        val epoch = buffer.getLong()
        val value = buffer.getFloat()
        (epoch, value)
      }
}


trait SparkStreamingProcessor extends SparkProcessor {
  def setupStreaming(args: Array[String]) = {
    val (properties, target, logLevel, sc, inputNatsStreaming, inputSubject, outputSubject, clusterId, outputNatsStreaming, natsUrl) = setup(args)

    val streamingDuration = scala.util.Properties.envOrElse("STREAMING_DURATION", "2000").toInt
    println("STREAMING_DURATION = " + streamingDuration)

    val ssc = new StreamingContext(sc, new Duration(streamingDuration));
//    ssc.checkpoint("/spark/storage")

    (properties, target, logLevel, sc, ssc, inputNatsStreaming, inputSubject, outputSubject, clusterId, outputNatsStreaming, natsUrl, streamingDuration)
  }
} 
Example 188
Source File: SparkTemperatureProcessor.scala    From smart-meter   with MIT License 5 votes vote down vote up
package com.logimethods.nats.connector.spark.app

import java.util.Properties;
import java.io.File
import java.io.Serializable

import org.apache.spark.SparkConf
import org.apache.spark.SparkContext
import org.apache.spark.storage.StorageLevel;
import org.apache.spark.streaming._
import com.datastax.spark.connector.streaming._
import com.datastax.spark.connector.SomeColumns

import io.nats.client.ConnectionFactory._
import java.nio.ByteBuffer

import org.apache.log4j.{Level, LogManager, PropertyConfigurator}

import com.logimethods.connector.nats.to_spark._
import com.logimethods.scala.connector.spark.to_nats._

import java.util.function._

import java.time.{LocalDateTime, ZoneOffset}

object SparkTemperatureProcessor extends App with SparkStreamingProcessor {
  val log = LogManager.getRootLogger
  log.setLevel(Level.WARN)

  val (properties, target, logLevel, sc, ssc, inputNatsStreaming, inputSubject, outputSubject, clusterId, outputNatsStreaming, natsUrl, streamingDuration) =
    setupStreaming(args)

  // Temperatures //

  val temperatures =
    if (inputNatsStreaming) {
      NatsToSparkConnector
        .receiveFromNatsStreaming(classOf[Tuple2[Long,Float]], StorageLevel.MEMORY_ONLY, clusterId)
        .withNatsURL(natsUrl)
        .withSubjects(inputSubject)
        .withDataDecoder(dataDecoder)
        .asStreamOf(ssc)
    } else {
      NatsToSparkConnector
        .receiveFromNats(classOf[Tuple2[Long,Float]], StorageLevel.MEMORY_ONLY)
        .withProperties(properties)
        .withSubjects(inputSubject)
        .withDataDecoder(dataDecoder)
        .asStreamOf(ssc)
    }

  // Ideally, should be the AVG
  val singleTemperature = temperatures.reduceByKey(Math.max(_,_))

  if (logLevel.contains("TEMPERATURE")) {
    singleTemperature.print()
  }

  singleTemperature.saveToCassandra("smartmeter", "temperature")

  val temperatureReport = singleTemperature.map({case (epoch, temperature) => (s"""{"epoch": $epoch, "temperature": $temperature}""") })
  SparkToNatsConnectorPool.newPool()
                      .withProperties(properties)
                      .withSubjects(outputSubject) // "smartmeter.extract.temperature"
                      .publishToNats(temperatureReport)

  // Start //
  ssc.start();

  ssc.awaitTermination()
} 
Example 189
Source File: ValueProviderTest.scala    From smart-meter   with MIT License 5 votes vote down vote up
package com.logimethods.smartmeter.generate

import com.logimethods.smartmeter.generate._
import org.scalatest._
import java.time._

class ValueProviderTest extends FunSuite {
  import java.nio.ByteBuffer

  val dataDecoder = new java.util.function.Function[Array[Byte],Tuple2[Long,Float]] {
    override def apply(bytes: Array[Byte]):Tuple2[Long,Float] = {
      val buffer = ByteBuffer.wrap(bytes);
      val epoch = buffer.getLong()
      val voltage = buffer.getFloat()
      (epoch, voltage)  
    }
  }

  test("encodePayload(date, value)") {
    val date = LocalDateTime.now()
    val value = 12345.6789f
    
    val bytes = new ConsumerInterpolatedVoltageProvider(1, 100, 5000, 0, 12).encodePayload(date, value)
    // print(new String(ByteBuffer.wrap(bytes).array()))
    
    val tuple = dataDecoder.apply(bytes)    
    // print(tuple)
    
    assert(date.atOffset(ZoneOffset.MIN).toEpochSecond() == tuple._1)
    assert(date.withNano(0) == LocalDateTime.ofEpochSecond(tuple._1, 0, ZoneOffset.MIN))
    assert(value == tuple._2)
  }

  test("computeNbOfElements(usersPerSec: Double)") {
      for(i <- 1 to 21){
         val nb = math.pow(2, i)
         val (lineNb, transformerNb, usagePointNb) = ProviderUtil.computeNbOfElements(nb)
         assert(lineNb > 0)
         assert(transformerNb > 0)
         assert(usagePointNb > 0)
         
         assert(usagePointNb >= transformerNb)
         assert(transformerNb >= lineNb)
      }    
  }
  

} 
Example 190
Source File: ModelMatrixEncoding.scala    From modelmatrix   with Apache License 2.0 5 votes vote down vote up
package com.collective.modelmatrix

import java.nio.ByteBuffer

import scodec.bits.ByteVector



private[modelmatrix] trait Encode[T] {
  def encode(obj: T): ByteVector
}

private[modelmatrix] object Encode {
  def apply[T](f: T => ByteVector): Encode[T] = new Encode[T] {
    def encode(obj: T): ByteVector = f(obj)
  }

  implicit val EncodeShort: Encode[Short] =
    Encode((short: Short) => ByteVector(ByteBuffer.allocate(2).putShort(short).array()))
  implicit val EncodeInt: Encode[Int] =
    Encode((integer: Int) => ByteVector(ByteBuffer.allocate(4).putInt(integer).array()))
  implicit val EncodeLong: Encode[Long] =
    Encode((long: Long) => ByteVector(ByteBuffer.allocate(8).putLong(long).array()))
  implicit val EncodeDouble: Encode[Double] =
    Encode((double: Double) => ByteVector(ByteBuffer.allocate(8).putDouble(double).array()))
  implicit val EncodeString: Encode[String] =
    Encode((string: String) => ByteVector(string.getBytes))

}

private[modelmatrix] trait Decode[T] {
  def decode(bytes: ByteVector): T
}

private[modelmatrix] object Decode {
  def apply[T](f: ByteVector => T): Decode[T] = new Decode[T] {
    def decode(bytes: ByteVector): T = f(bytes)
  }

  implicit val DecodeShort: Decode[Short] =
    Decode(_.toByteBuffer.getShort)
  implicit val DecodeInt: Decode[Int] =
    Decode(_.toByteBuffer.getInt)
  implicit val DecodeLong: Decode[Long] =
    Decode(_.toByteBuffer.getLong)
  implicit val DecodeDouble: Decode[Double] =
    Decode(_.toByteBuffer.getDouble)
  implicit val DecodeString: Decode[String] =
    Decode(b => new String(b.toArray))

}

class ModelMatrixEncoding {

  def encode[T: Encode](obj: T): ByteVector = implicitly[Encode[T]].encode(obj)

  def decode[T: Decode](obj: ByteVector): T = implicitly[Decode[T]].decode(obj)

}

object ModelMatrixEncoding extends ModelMatrixEncoding 
Example 191
Source File: MessageBuffer.scala    From scala-loci   with Apache License 2.0 5 votes vote down vote up
package loci

import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets

import scala.annotation.compileTimeOnly
import scala.collection.mutable

final class MessageBuffer private (val backingArray: Array[Byte])
    extends mutable.IndexedSeq[Byte] {
  @compileTimeOnly("`backingArrayBuffer` only available in JS")
  def backingArrayBuffer: Any = ???

  @inline def length: Int = backingArray.length

  @inline def apply(index: Int) = {
    if (index < 0  || index >= length)
      throw new IndexOutOfBoundsException(s"index $index")

    backingArray(index)
  }

  @inline def update(index: Int, element: Byte) = {
    if (index < 0  || index >= length)
      throw new IndexOutOfBoundsException(s"index $index")

    backingArray(index) = element
  }

  @inline def update(offset: Int, buffer: MessageBuffer, bufferOffset: Int, count: Int) = {
    if (offset < 0 || bufferOffset < 0 || count < 0 ||
        offset > length - count || bufferOffset > buffer.length - count)
      throw new IndexOutOfBoundsException(
        s"offset $offset, length $length, " +
        s"buffer offset ${bufferOffset}, buffer length ${buffer.length}, count $count")

    System.arraycopy(buffer.backingArray, bufferOffset, backingArray, offset, count)
  }

  @inline def concat(buffer: MessageBuffer): MessageBuffer = {
    val array = new Array[Byte](length + buffer.length)
    System.arraycopy(backingArray, 0, array, 0, length)
    System.arraycopy(buffer.backingArray, 0, array, length, buffer.length)
    new MessageBuffer(array)
  }

  @inline def copy(offset: Int, count: Int): MessageBuffer = {
    if (offset < 0 || count < 0 || offset > length - count)
      throw new IndexOutOfBoundsException(s"offset $offset, count $count, length $length")

    val array = new Array[Byte](count)
    System.arraycopy(backingArray, offset, array, 0, count)
    new MessageBuffer(array)
  }

  @inline def decodeString(offset: Int, count: Int): String =
    new String(backingArray, offset, count, StandardCharsets.UTF_8)

  @inline def decodeString: String =
    decodeString(0, length)

  @inline def asByteBuffer: ByteBuffer =
    ByteBuffer wrap backingArray

  override def toString: String =
    MessageBufferEncoding.byteBufferToString(asByteBuffer, 0, length, fatal = true) getOrElse
      MessageBufferEncoding.messageBufferToHexString(this)
}

object MessageBuffer {
  def empty: MessageBuffer = new MessageBuffer(Array.emptyByteArray)

  def allocate(length: Int): MessageBuffer = new MessageBuffer(new Array(length))

  def encodeString(string: String): MessageBuffer =
    new MessageBuffer(string getBytes StandardCharsets.UTF_8)

  def wrapByteBuffer(buffer: ByteBuffer): MessageBuffer =
    if (!buffer.hasArray) {
      val duplicate = buffer.duplicate
      duplicate.position(0)
      duplicate.limit(buffer.capacity)
      val array = new Array[Byte](duplicate.remaining)
      duplicate.get(array)
      new MessageBuffer(array)
    }
    else
      new MessageBuffer(buffer.array)

  def wrapArray(array: Array[Byte]): MessageBuffer =
    new MessageBuffer(array)

  @compileTimeOnly("`wrapArrayBuffer` only available in JS")
  def wrapArrayBuffer(arrayBuffer: Any): MessageBuffer = ???
} 
Example 192
Source File: MessageBufferEncoding.scala    From scala-loci   with Apache License 2.0 5 votes vote down vote up
package loci

import java.nio.charset.{CharacterCodingException, CodingErrorAction, StandardCharsets}
import java.nio.{ByteBuffer, CharBuffer}

import scala.util.{Failure, Success, Try}

private object MessageBufferEncoding {
  locally(MessageBufferEncoding)

  def byteBufferToString(byteBuffer: ByteBuffer, offset: Int, count: Int, fatal: Boolean): Try[String] = {
    val decoder = StandardCharsets.UTF_8.newDecoder()

    if (!fatal)
      decoder
        .onMalformedInput(CodingErrorAction.REPLACE)
        .onUnmappableCharacter(CodingErrorAction.REPLACE)

    val size = (count * decoder.maxCharsPerByte.toDouble).toInt
    val array = new Array[Char](size)
    val charBuffer = CharBuffer.wrap(array)

    byteBuffer.position(offset)
    byteBuffer.limit(offset + count)

    try {
      var result = decoder.decode(byteBuffer, charBuffer, true)

      if (!result.isUnderflow)
        result.throwException()
      result = decoder.flush(charBuffer)
      if (!result.isUnderflow)
        result.throwException()

      Success(new String(array, 0, charBuffer.position()))
    }
    catch {
      case exception: CharacterCodingException =>
        Failure(exception)
    }
  }

  def stringToByteBuffer(string: String)(allocateByteBuffer: Int => ByteBuffer): ByteBuffer = {
    val encoder = StandardCharsets.UTF_8.newEncoder()
      .onMalformedInput(CodingErrorAction.REPLACE)
      .onUnmappableCharacter(CodingErrorAction.REPLACE)

    val size = (string.length * encoder.maxBytesPerChar.toDouble).toInt
    val byteBuffer = allocateByteBuffer(size)
    val charBuffer = CharBuffer.wrap(string)

    byteBuffer.position(0)
    byteBuffer.limit(size)

    var result = encoder.encode(charBuffer, byteBuffer, true)

    if (!result.isUnderflow)
      result.throwException()
    result = encoder.flush(byteBuffer)
    if (!result.isUnderflow)
      result.throwException()

    byteBuffer
  }

  private val hex = Array(
    '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F')

  def messageBufferToHexString(buffer: MessageBuffer): String = {
    val result = new Array[Char](3 * buffer.length)
    var i = 0
    var j = 0
    while (i < buffer.length) {
      result(j) = hex((buffer(i) & 0xF0) >> 4)
      j += 1
      result(j) = hex(buffer(i) & 0x0F)
      j += 1
      result(j) = ' '
      j += 1
      i += 1
    }

    if (result.isEmpty)
      ""
    else
      new String(result, 0, result.length - 1)
  }
} 
Example 193
Source File: jsoniterScala.scala    From scala-loci   with Apache License 2.0 5 votes vote down vote up
package loci
package serializer

import scala.util.Try
import java.nio.ByteBuffer

import com.github.plokhotnyuk.jsoniter_scala.core._
import loci.transmitter.Serializable

object jsoniterScala {
  implicit def jsoniteScalaBasedSerializable[T]
  (implicit codec: JsonValueCodec[T]): Serializable[T] = new Serializable[T] {
    def serialize(value: T): MessageBuffer = {
      val bytes = writeToArray(value)
      MessageBuffer.wrapByteBuffer(ByteBuffer.wrap(bytes))
    }
    def deserialize(value: MessageBuffer): Try[T] =
      Try {readFromByteBuffer(value.asByteBuffer)}
  }
} 
Example 194
Source File: HasTypeCodec.scala    From troy   with Apache License 2.0 5 votes vote down vote up
package troy
package driver.codecs

import java.net.InetAddress
import java.nio.ByteBuffer
import java.util.{ Date, UUID }

import com.datastax.driver.core.{ LocalDate, TypeCodec }
import troy.driver.{ CassandraDataType => CT }

import scala.annotation.implicitNotFound

case class HasTypeCodec[S, C <: CT](typeCodec: TypeCodec[S]) extends AnyVal

object HasTypeCodec {
  import TypeCodec._

  implicit val BooleanHasCodecAsBoolean = HasTypeCodec[java.lang.Boolean, CT.Boolean](cboolean)
  implicit val TinyIntHasCodecAsByte = HasTypeCodec[java.lang.Byte, CT.TinyInt](tinyInt)
  implicit val SmallIntHasCodecAsShort = HasTypeCodec[java.lang.Short, CT.SmallInt](smallInt)
  implicit val IntHasCodecAsInteger = HasTypeCodec[java.lang.Integer, CT.Int](cint)
  implicit val BigIntHasCodecAsLong = HasTypeCodec[java.lang.Long, CT.BigInt](bigint)
  implicit val CounterHasCodecAsLong = HasTypeCodec[java.lang.Long, CT.Counter](counter)
  implicit val FloatHasCodecAsFloat = HasTypeCodec[java.lang.Float, CT.Float](cfloat)
  implicit val DoubleHasCodecAsDouble = HasTypeCodec[java.lang.Double, CT.Double](cdouble)
  implicit val VarIntHasCodecAsBigInteger = HasTypeCodec[java.math.BigInteger, CT.VarInt](varint)
  implicit val DecimalHasCodecAsBigDecimal = HasTypeCodec[java.math.BigDecimal, CT.Decimal](decimal)
  implicit val AsciiHasCodecAsString = HasTypeCodec[String, CT.Ascii](ascii)
  implicit val VarCharHasCodecAsString = HasTypeCodec[String, CT.VarChar](varchar)
  implicit val TextHasCodecAsString = HasTypeCodec[String, CT.Text](varchar)
  implicit val BlobHasCodecAsByteBuffer = HasTypeCodec[ByteBuffer, CT.Blob](blob)
  implicit val DateHasCodecAsLocalDate = HasTypeCodec[LocalDate, CT.Date](date)
  implicit val TimeHasCodecAsLong = HasTypeCodec[java.lang.Long, CT.Time](time)
  implicit val TimestampHasCodecAsDate = HasTypeCodec[Date, CT.Timestamp](timestamp)
  implicit val UuidHasCodecAsUUID = HasTypeCodec[UUID, CT.Uuid](uuid)
  implicit val TimeUuidHasCodecAsUUID = HasTypeCodec[UUID, CT.TimeUuid](timeUUID)
  implicit val InetHasCodecAsInetAddress = HasTypeCodec[InetAddress, CT.Inet](inet)
} 
Example 195
Source File: FakeAbstractGettableByIndexData.scala    From troy   with Apache License 2.0 5 votes vote down vote up
package com.datastax.driver.core

import java.net.InetAddress
import java.nio.ByteBuffer
import java.util.Date

class FakeAbstractGettableByIndexData(pv: ProtocolVersion, data: (ByteBuffer, DataType)*) extends AbstractGettableByIndexData(pv) {
  override def getType(i: Int): DataType = data(i)._2

  override def getValue(i: Int): ByteBuffer = data(i)._1

  override def getName(i: Int): String = ""

  override def getCodecRegistry: CodecRegistry = ???
}
object FakeAbstractGettableByIndexData {
  def ascii(value: String, pv: ProtocolVersion = ProtocolVersion.V4) = new FakeAbstractGettableByIndexData(
    pv, TypeCodec.ascii.serialize(value, pv) -> DataType.ascii()
  )

  def timestamp(value: Date, pv: ProtocolVersion = ProtocolVersion.V4) = new FakeAbstractGettableByIndexData(
    pv, TypeCodec.timestamp.serialize(value, pv) -> DataType.timestamp()
  )

  def decimal(value: java.math.BigDecimal, pv: ProtocolVersion = ProtocolVersion.V4) = new FakeAbstractGettableByIndexData(
    pv, TypeCodec.decimal.serialize(value, pv) -> DataType.decimal()
  )

  def inet(value: InetAddress, pv: ProtocolVersion = ProtocolVersion.V4) = new FakeAbstractGettableByIndexData(
    pv, TypeCodec.inet.serialize(value, pv) -> DataType.inet()
  )

  def date(value: LocalDate, pv: ProtocolVersion = ProtocolVersion.V4) = new FakeAbstractGettableByIndexData(
    pv, TypeCodec.date.serialize(value, pv) -> DataType.date()
  )
} 
Example 196
Source File: ReadingWritingData.scala    From Spark-RSVD   with Apache License 2.0 5 votes vote down vote up
package com.criteo.rsvd

import java.nio.ByteBuffer

import com.esotericsoftware.kryo.Kryo
import com.typesafe.scalalogging.slf4j.StrictLogging
import de.javakaffee.kryoserializers.UnmodifiableCollectionsSerializer
import org.apache.hadoop.fs.{FileSystem, Path}
import org.apache.hadoop.io.{BytesWritable, NullWritable}
import org.apache.spark.mllib.linalg.distributed.MatrixEntry
import org.apache.spark.rdd.RDD
import org.apache.spark.serializer.{KryoRegistrator, KryoSerializer}
import org.apache.spark.{SparkConf, SparkContext}

import scala.reflect.ClassTag

object ReadingWritingData extends StrictLogging {

  def getInputDataSizeMB(inputPathPattern: String, sc: SparkContext): Int = {
    val fs = FileSystem.get(sc.hadoopConfiguration)
    val path = new Path(inputPathPattern)
    (fs.globStatus(path).map(f => f.getLen).sum / 1024 / 1024).toInt
  }

  def loadMatrixEntries(inputPath: String,
                        singlePartitionSizeMB: Int,
                        sc: SparkContext): RDD[MatrixEntry] = {

    logger.info(s"Input matrix path: $inputPath")
    val inputDataSizeMB = getInputDataSizeMB(inputPath + "
  def makeRddFromKryoFile[T: ClassTag](
      sc: SparkContext,
      path: String,
      minPartitionsOpt: Option[Int] = None): RDD[T] = {
    val minPartitions = minPartitionsOpt.getOrElse(sc.defaultMinPartitions)
    val serializer = new KryoSerializer(sc.getConf)
    sc.sequenceFile(path,
                    classOf[NullWritable],
                    classOf[BytesWritable],
                    minPartitions)
      .mapPartitions { it =>
        val instance = serializer.newInstance()
        it.flatMap {
          case (_, v) =>
            instance.deserialize[Array[T]](ByteBuffer.wrap(v.getBytes))
        }
      }
  }

  object RandomizedSVDKryoRegistrator extends KryoRegistrator {

    def registerClasses(kryo: Kryo): Unit = {
      UnmodifiableCollectionsSerializer.registerSerializers(kryo)
      kryo.register(classOf[MatrixEntry])
      kryo.register(classOf[Array[MatrixEntry]])
    }
  }

  def appendBasicRegistratorToSparkConf(sparkConf: SparkConf): SparkConf =
    appendRegistratorToSparkConf(sparkConf,
                                 RandomizedSVDKryoRegistrator.getClass.getName)

  def appendRegistratorToSparkConf(sparkConf: SparkConf,
                                   registratorName: String): SparkConf = {
    val oldValue = sparkConf.get("spark.kryo.registrator", "")
    if (oldValue == "") {
      sparkConf.set("spark.kryo.registrator", registratorName)
    } else {
      sparkConf.set("spark.kryo.registrator", oldValue + "," + registratorName)
    }
  }

} 
Example 197
Source File: Serdes.scala    From tamer   with MIT License 5 votes vote down vote up
package tamer

import java.io.ByteArrayOutputStream
import java.nio.ByteBuffer

import com.sksamuel.avro4s._
import org.apache.avro.Schema
import tamer.registry._
import zio.{RIO, Task}
import zio.kafka.client.serde.{Deserializer, Serializer}

sealed trait Serde[A] extends Any {
  def isKey: Boolean
  def schema: Schema
  def deserializer: Deserializer[Registry with Topic, A]
  def serializer: Serializer[Registry with Topic, A]
  final def serde: ZSerde[Registry with Topic, A] = ZSerde(deserializer)(serializer)
}

object Serde {
  private[this] final val Magic: Byte = 0x0
  private[this] final val intByteSize = 4

  final def apply[A <: Product: Decoder: Encoder: SchemaFor](isKey: Boolean = false) =
    new RecordSerde[A](isKey, SchemaFor[A].schema(DefaultFieldMapper))

  final class RecordSerde[A: Decoder: Encoder](override final val isKey: Boolean, override final val schema: Schema) extends Serde[A] {
    private[this] def subject(topic: String): String = s"$topic-${if (isKey) "key" else "value"}"
    override final val deserializer: Deserializer[Registry with Topic, A] = Deserializer.byteArray.mapM { ba =>
      val buffer = ByteBuffer.wrap(ba)
      if (buffer.get() != Magic) RIO.fail(SerializationError("Unknown magic byte!"))
      else {
        val id = buffer.getInt()
        for {
          env <- RIO.environment[Registry]
          _   <- env.registry.verifySchema(id, schema)
          res <- RIO.fromTry {
            val length  = buffer.limit() - 1 - intByteSize
            val payload = new Array[Byte](length)
            buffer.get(payload, 0, length)
            AvroInputStream.binary[A].from(payload).build(schema).tryIterator.next
          }
        } yield res
      }
    }
    override final val serializer: Serializer[Registry with Topic, A] = Serializer.byteArray.contramapM { a =>
      for {
        env <- RIO.environment[Registry with Topic]
        id  <- env.registry.getOrRegisterId(subject(env.topic), schema)
        arr <- Task {
          val baos = new ByteArrayOutputStream
          baos.write(Magic.toInt)
          baos.write(ByteBuffer.allocate(intByteSize).putInt(id).array())
          val ser = AvroOutputStream.binary[A].to(baos).build(schema)
          ser.write(a)
          ser.close()
          baos.toByteArray
        }
      } yield arr
    }
  }
} 
Example 198
Source File: TrafficMonitorThread.scala    From shadowsocksr-android   with GNU General Public License v3.0 5 votes vote down vote up
package com.github.shadowsocks.utils

import java.io.{File, IOException}
import java.nio.{ByteBuffer, ByteOrder}
import java.util.concurrent.Executors

import android.content.Context
import android.net.{LocalServerSocket, LocalSocket, LocalSocketAddress}
import android.util.Log

class TrafficMonitorThread(context: Context) extends Thread {

  val TAG = "TrafficMonitorThread"
  lazy val PATH = context.getApplicationInfo.dataDir + "/stat_path"

  @volatile var serverSocket: LocalServerSocket = null
  @volatile var isRunning: Boolean = true

  def closeServerSocket() {
    if (serverSocket != null) {
      try {
        serverSocket.close()
      } catch {
        case _: Exception => // ignore
      }
      serverSocket = null
      }
  }

  def stopThread() {
    isRunning = false
    closeServerSocket()
  }

  override def run() {

    try {
      new File(PATH).delete()
    } catch {
      case _: Exception => // ignore
    }

    try {
      val localSocket = new LocalSocket
      localSocket.bind(new LocalSocketAddress(PATH, LocalSocketAddress.Namespace.FILESYSTEM))
      serverSocket = new LocalServerSocket(localSocket.getFileDescriptor)
    } catch {
      case e: IOException =>
        Log.e(TAG, "unable to bind", e)
        return
    }

    val pool = Executors.newFixedThreadPool(1)

    while (isRunning) {
      try {
        val socket = serverSocket.accept()

        pool.execute(() => {
          try {
            val input = socket.getInputStream
            val output = socket.getOutputStream

            val buffer = new Array[Byte](16)
            if (input.read(buffer) != 16) throw new IOException("Unexpected traffic stat length")
            val stat = ByteBuffer.wrap(buffer).order(ByteOrder.LITTLE_ENDIAN)
            TrafficMonitor.update(stat.getLong(0), stat.getLong(8))

            output.write(0)

            input.close()
            output.close()

          } catch {
            case e: Exception =>
              Log.e(TAG, "Error when recv traffic stat", e)
          }

          // close socket
          try {
            socket.close()
          } catch {
            case _: Exception => // ignore
          }

        })
      } catch {
        case e: IOException =>
          Log.e(TAG, "Error when accept socket", e)
          return
      }
    }
  }
} 
Example 199
Source File: Customer.scala    From fusion-data   with Apache License 2.0 5 votes vote down vote up
package kafkasample.demo

import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets
import java.util

import helloscala.common.util.StringUtils
import org.apache.kafka.common.serialization.Serializer

case class Customer(customerId: Int, customerName: String) {}

class CustomerSerializer extends Serializer[Customer] {
  private val EMPTY_NAME = Array[Byte]()

  override def configure(configs: util.Map[String, _], isKey: Boolean): Unit = {}

  override def serialize(topic: String, data: Customer): Array[Byte] =
    if (data eq null) {
      null
    } else {
      var nameLen = 0
      var nameBytes = EMPTY_NAME
      if (StringUtils.isNoneBlank(data.customerName)) {
        nameLen = data.customerName.length
        nameBytes = data.customerName.getBytes(StandardCharsets.UTF_8)
      }

      val buf = ByteBuffer.allocate(4 + 4 + nameLen)
      buf.putInt(data.customerId)
      buf.putInt(nameLen)
      buf.put(nameBytes)

      buf.array()
    }

  override def close(): Unit = ???
} 
Example 200
Source File: AvroMessageConverterTest.scala    From stream-reactor   with Apache License 2.0 5 votes vote down vote up
package com.datamountaineer.streamreactor.connect.sink.converters

import java.nio.ByteBuffer
import java.util.UUID

import com.datamountaineer.streamreactor.connect.jms.config.{JMSConfig, JMSSettings}
import com.datamountaineer.streamreactor.connect.jms.sink.converters.AvroMessageConverter
import com.datamountaineer.streamreactor.connect.sink.AvroDeserializer
import com.datamountaineer.streamreactor.connect.{TestBase, Using}
import io.confluent.connect.avro.AvroData
import javax.jms.BytesMessage
import org.apache.activemq.ActiveMQConnectionFactory
import org.apache.avro.generic.GenericData
import org.apache.avro.util.Utf8
import org.scalatest.BeforeAndAfterAll
import org.scalatest.matchers.should.Matchers
import org.scalatest.wordspec.AnyWordSpec

import scala.collection.JavaConverters._
import scala.reflect.io.Path

class AvroMessageConverterTest extends AnyWordSpec with Matchers with Using with TestBase with BeforeAndAfterAll {
  val converter = new AvroMessageConverter()
  private lazy val avroData = new AvroData(128)
  val kafkaTopic1 = s"kafka-${UUID.randomUUID().toString}"
  val topicName = UUID.randomUUID().toString
  val queueName = UUID.randomUUID().toString

  val kcqlT = getKCQL(topicName, kafkaTopic1, "TOPIC")
  val kcqlQ = getKCQL(queueName, kafkaTopic1, "QUEUE")

  val props = getProps(s"$kcqlQ;$kcqlT", JMS_URL)
  val config = JMSConfig(props.asJava)
  val settings = JMSSettings(config, true)
  val setting = settings.settings.head

  override def afterAll(): Unit = {
    Path(AVRO_FILE).delete()
  }

  "AvroMessageConverter" should {
    "create a BytesMessage with avro payload" in {
      val connectionFactory = new ActiveMQConnectionFactory("vm://localhost?broker.persistent=false")

      using(connectionFactory.createConnection()) { connection =>
        using(connection.createSession(false, 1)) { session =>

          val record = getSinkRecords(kafkaTopic1).head
          val msg = converter.convert(record, session, setting)._2.asInstanceOf[BytesMessage]

          Option(msg).isDefined shouldBe true

          msg.reset()

          val size = msg.getBodyLength

          size > 0 shouldBe true
          val data = new Array[Byte](size.toInt)
          msg.readBytes(data)

          val avroRecord = AvroDeserializer(data, avroData.fromConnectSchema(record.valueSchema()))
          avroRecord.get("int8") shouldBe 12.toByte
          avroRecord.get("int16") shouldBe 12.toShort
          avroRecord.get("int32") shouldBe 12
          avroRecord.get("int64") shouldBe 12L
          avroRecord.get("float32") shouldBe 12.2f
          avroRecord.get("float64") shouldBe 12.2
          avroRecord.get("boolean") shouldBe true
          avroRecord.get("string").toString shouldBe "foo"
          avroRecord.get("bytes").asInstanceOf[ByteBuffer].array() shouldBe "foo".getBytes()
          val array = avroRecord.get("array").asInstanceOf[GenericData.Array[Utf8]]
          val iter = array.iterator()
          new Iterator[String] {
            override def hasNext: Boolean = iter.hasNext

            override def next(): String = iter.next().toString
          }.toSeq shouldBe Seq("a", "b", "c")
          val map = avroRecord.get("map").asInstanceOf[java.util.Map[Utf8, Int]].asScala
          map.size shouldBe 1
          map.keys.head.toString shouldBe "field"
          map.get(map.keys.head) shouldBe Some(1)

          val iterRecord = avroRecord.get("mapNonStringKeys").asInstanceOf[GenericData.Array[GenericData.Record]].iterator()
          iterRecord.hasNext shouldBe true
          val r = iterRecord.next()
          r.get("key") shouldBe 1
          r.get("value") shouldBe 1
        }
      }
    }
  }
}