<commit_before>/**
 * Copyright (c) 2015 Basho Technologies, Inc.
 *
 * This file is provided to you under the Apache License,
 * Version 2.0 (the "License"); you may not use this file
 * except in compliance with the License.  You may obtain
 * a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package com.basho.spark.connector.query

import com.basho.riak.client.api.cap.Quorum
import org.apache.spark.Logging

import scala.collection.JavaConversions._

import com.basho.riak.client.api.RiakClient
import com.basho.riak.client.api.commands.kv.{FetchValue, MultiFetch}
import com.basho.riak.client.core.query.{RiakObject, Location}

import scala.collection.mutable.ArrayBuffer

class DataQueryingIterator(query: Query[_], riakSession: RiakClient, minConnectionsPerNode: Int)
  extends Iterator[(Location, RiakObject)] with Logging {

  type ResultT = (Location, RiakObject)

  // The following 2 variables are useful only for the debugging purposes
  private val dataBuffer: ArrayBuffer[ResultT] = new ArrayBuffer[ResultT](query.readConf.fetchSize)
  private var bufferIndex = 0

  private var isThereANextValue: Any = None
  private var nextToken: Option[_] = None

  private var _iterator: Option[Iterator[ResultT]] = None

  protected[this] def prefetchIfNeeded(): Boolean = {
    // scalastyle:off return
    if (this._iterator.isDefined) {
      if(this._iterator.get.hasNext){
        // Nothing to do, we still have at least one result
        logDebug(s"prefetch is not required, at least one pre-fetched value [actually ${dataBuffer.size-bufferIndex} from ${dataBuffer.size}] is in the buffer")
        return false
      }else if(nextToken.isEmpty){
        // Nothing to do, we already got all the data
        logDebug("prefetch is not required, all data has been processed")
        return false
      }
    }
    // scalastyle:on return

    logTrace(s"Performing 2i query(token=$nextToken)")

    val r = query.nextLocationBulk(nextToken, riakSession )
    logDebug(s"2i query(token=$nextToken) returns:\n  token: ${r._1}\n  locations: ${r._2}")
    nextToken = r._1

    dataBuffer.clear()
    bufferIndex = 0

    if(r._2.isEmpty){
      /**
       * It is Absolutely possible situation, for instance:
       *     in case when the last data page will be returned as a result of  2i continuation query and
       *     this page will be fully filled with data then the valid continuation token wile be also returned (it will be not null),
       *     therefore additional/subsequent data fetch request will be required.
       *     As a result of such call the empty locations list and Null continuation token will be returned
       */
      logDebug("prefetch is not required, all data was processed (location list is empty)")
      _iterator = Some(Iterator.empty)
    } else {

      /**
       * To be 100% sure that massive fetch doesn't lead to the connection pool starvation,
       * fetch will be performed by the smaller chunk of keys.
       *
       * Ideally the chunk size should be equal to min number of connections to the RiakNode
       */
      val itChunkedLocations: Iterator[Iterable[Location]] = r._2.grouped(minConnectionsPerNode)
      DataQueryingIterator.fetchData(riakSession, itChunkedLocations, dataBuffer)

      logDebug(s"Next data buffer was fetched:\n" +
        s"\tnextToken: $nextToken\n" +
        s"\tbuffer: $dataBuffer")

      _iterator = Some(dataBuffer.iterator)
    }
    true
  }

  override def hasNext: Boolean =
    isThereANextValue match {
      case b: Boolean =>
        b
      case _ =>
        prefetchIfNeeded()
        val r= _iterator.get.hasNext
        isThereANextValue = r
        r
    }

  override def next(): (Location, RiakObject) = {
    if( !hasNext ){
      throw new NoSuchElementException("next on iterator")
    }

    bufferIndex += 1
    isThereANextValue = None<commit_msg>DataQueryingIterator  if was replaced by pattern matching<commit_after>/**
 * Copyright (c) 2015 Basho Technologies, Inc.
 *
 * This file is provided to you under the Apache License,
 * Version 2.0 (the "License"); you may not use this file
 * except in compliance with the License.  You may obtain
 * a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package com.basho.spark.connector.query

import com.basho.riak.client.api.cap.Quorum
import org.apache.spark.Logging

import scala.collection.JavaConversions._

import com.basho.riak.client.api.RiakClient
import com.basho.riak.client.api.commands.kv.{FetchValue, MultiFetch}
import com.basho.riak.client.core.query.{RiakObject, Location}

import scala.collection.mutable.ArrayBuffer

class DataQueryingIterator(query: Query[_], riakSession: RiakClient, minConnectionsPerNode: Int)
  extends Iterator[(Location, RiakObject)] with Logging {

  type ResultT = (Location, RiakObject)

  // The following 2 variables are useful only for the debugging purposes
  private val dataBuffer: ArrayBuffer[ResultT] = new ArrayBuffer[ResultT](query.readConf.fetchSize)
  private var bufferIndex = 0

  private var isThereANextValue: Any = None
  private var nextToken: Option[_] = None

  private var _iterator: Option[Iterator[ResultT]] = None

  protected[this] def prefetchIfNeeded(): Boolean = {
    // scalastyle:off return
    if (this._iterator.isDefined) {
      if(this._iterator.get.hasNext){
        // Nothing to do, we still have at least one result
        logDebug(s"prefetch is not required, at least one pre-fetched value [actually ${dataBuffer.size-bufferIndex} from ${dataBuffer.size}] is in the buffer")
        return false
      }else if(nextToken.isEmpty){
        // Nothing to do, we already got all the data
        logDebug("prefetch is not required, all data has been processed")
        return false
      }
    }
    // scalastyle:on return

    logTrace(s"Performing 2i query(token=$nextToken)")

    val r = query.nextLocationBulk(nextToken, riakSession )
    logDebug(s"2i query(token=$nextToken) returns:\n  token: ${r._1}\n  locations: ${r._2}")
    nextToken = r._1

    dataBuffer.clear()
    bufferIndex = 0

    r match {
      case (_, Nil) =>
        /**
         * It is Absolutely possible situation, for instance:
         *     in case when the last data page will be returned as a result of  2i continuation query and
         *     this page will be fully filled with data then the valid continuation token wile be also returned (it will be not null),
         *     therefore additional/subsequent data fetch request will be required.
         *     As a result of such call the empty locations list and Null continuation token will be returned
         */
        logDebug("prefetch is not required, all data was processed (location list is empty)")
        _iterator = Some(Iterator.empty)
      case (_, locations: Iterable[Location]) =>
        /**
         * To be 100% sure that massive fetch doesn't lead to the connection pool starvation,
         * fetch will be performed by the smaller chunks of keys.
         *
         * Ideally the chunk size should be equal to min number of connections to the RiakNode
         */
        val itChunkedLocations = locations.grouped(minConnectionsPerNode)
        DataQueryingIterator.fetchData(riakSession, itChunkedLocations, dataBuffer)

        logDebug(s"Next data buffer was fetched:\n" +
          s"\tnextToken: $nextToken\n" +
          s"\tbuffer: $dataBuffer")

        _iterator = Some(dataBuffer.iterator)
    }
    true
  }

  override def hasNext: Boolean =
    isThereANextValue match {
      case b: Boolean =>
        b
      case _ =>
        prefetchIfNeeded()
        val r= _iterator.get.hasNext
        isThereANextValue = r
        r
    }

  override def next(): (Location, RiakObject) = {
    if( !hasNext ){
      throw new NoSuchElementException("next on iterator")
    }

    bufferIndex += 1
    isThereANextValue = None
{'fraction_non_alphanumeric': 0.0840914392348962, 'fraction_numerical': 0.0051317937951947746, 'mean_word_length': 3.077508321445554, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5048758', 'n_tokens_mistral': 2471, 'n_tokens_neox': 2370, 'n_words': 851}