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

package unicorn.graph

import unicorn.json._
import unicorn.unibase.{Bucket, UniBase}
import unicorn.util.Logging

/** Abstract Unibase graph visitor. A default implementation of edges returns
  * an adjacency list by checking $.graph in the document.  The field $.graph
  * is an object (usually stored in a separate column family)
  * of which fields are relationships and values are target node IDs.
  * Each target node can have any data associate with.
  * For example,
  *
  * {{{
  *   {
  *     "graph": {
  *       "work with": {
  *         "Tom": {  // The field name is the string representation of _id, which could be a complecited JsValue.
  *           "_id": "Tom",
  *           "data": {
  *             "project": "Call of Duty",
  *             "since": 2015
  *           }
  *         }
  *       },
  *       "report to": {
  *         "Bob": {
  *           "_id": "Bob",
  *           "data": {
  *             "since": 2010
  *           }
  *         }
  *       }
  *     }
  *   }
  * }}}
  *
  * @author <NAME>
  */
class UniBaseVisitor(bucket: Bucket, maxHops: Int = 3) extends Visitor[JsValue, (String, JsValue)] with Logging {
  /** Cache of graph nodes in the database. */
  val cache = collection.mutable.Map[JsValue, JsObject]()

  /** Relationship of interest. Only neighbors with given relationship will be visited. */
  var relationships: Option[Set[String]] = None

  /** Map node to an index. */
  val nodes = scala.collection.mutable.Map[JsValue, Int]()
  /** Map edge to a weight. */
  val weights = scala.collection.mutable.Map[(Int, Int), Double]()

  override def visit(vertex: JsValue, edge: Edge[JsValue, (String, JsValue)], hops: Int): Unit = {
    if (!nodes.contains(vertex)) nodes(vertex) = nodes.size
    
    if (edge != null) {
      val weight = edge.data match {
        case Some((_, data: JsInt)) => data.value
        case Some((_, data: JsLong)) => data.value
        case Some((_, data: JsDouble)) => data.value
        case _ => 1.0
      }
      
      weights((nodes(edge.source), nodes(edge.target))) = weight
    }
  }<commit_msg>add visit hook<commit_after>/*******************************************************************************
 * (C) Copyright 2015 ADP, LLC.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *******************************************************************************/

package unicorn.graph

import unicorn.json._
import unicorn.unibase.{Bucket, UniBase}
import unicorn.util.Logging

/** Abstract Unibase graph visitor. A default implementation of edges returns
  * an adjacency list by checking $.graph in the document.  The field $.graph
  * is an object (usually stored in a separate column family)
  * of which fields are relationships and values are target node IDs.
  * Each target node can have any data associate with.
  * For example,
  *
  * {{{
  *   {
  *     "graph": {
  *       "work with": {
  *         "Tom": {  // The field name is the string representation of _id, which could be a complecited JsValue.
  *           "_id": "Tom",
  *           "data": {
  *             "project": "Call of Duty",
  *             "since": 2015
  *           }
  *         }
  *       },
  *       "report to": {
  *         "Bob": {
  *           "_id": "Bob",
  *           "data": {
  *             "since": 2010
  *           }
  *         }
  *       }
  *     }
  *   }
  * }}}
  *
  * @author <NAME>
  */
class UniBaseVisitor(bucket: Bucket, maxHops: Int = 3) extends Visitor[JsValue, (String, JsValue)] with Logging {
  /** Cache of graph nodes in the database. */
  val cache = collection.mutable.Map[JsValue, JsObject]()

  /** Relationship of interest. Only neighbors with given relationship will be visited. */
  var relationships: Option[Set[String]] = None

  /** Map node to an index. */
  val nodes = scala.collection.mutable.Map[JsValue, Int]()
  /** Map edge to a weight. */
  val weights = scala.collection.mutable.Map[(Int, Int), Double]()

  /** Extra action when visiting a vertex. This default implementation does nothing.
    * Applications needing customized visit action should override it. */
  val visitHooks = Seq[(JsValue, Edge[JsValue, (String, JsValue)], Int) => Unit]()

  override def visit(vertex: JsValue, edge: Edge[JsValue, (String, JsValue)], hops: Int): Unit = {
    if (!nodes.contains(vertex)) nodes(vertex) = nodes.size
    
    if (edge != null) {
      val weight = edge.data match {
        case Some((_, data: JsInt)) => data.value
        case Some((_, data: JsLong)) => data.value
        case Some((_, data: JsDouble)) => data.value
        case _ => 1.0
      }
      
      weights((nodes(edge.source), nodes(edge.target))) = weight
    }

    visitHooks.foreach { hook =>
      hook(vertex, edge, hops)
    }
  }
{'fraction_non_alphanumeric': 0.16719034231850505, 'fraction_numerical': 0.006284107822060526, 'mean_word_length': 2.658802177858439, 'pattern_counts': {'":': 24, '<': 5, '<?xml version=': 0, '>': 15, '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': '5932600', 'n_tokens_mistral': 1774, 'n_tokens_neox': 1686, 'n_words': 602}