| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * 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 org.apache.s2graph.core |
| |
| import org.apache.s2graph.core.mysqls.{Label, LabelMeta} |
| import org.apache.s2graph.core.types.{InnerVal, InnerValLikeWithTs, VertexId} |
| import org.apache.s2graph.core.utils.logger |
| |
| import scala.collection.mutable.{ArrayBuffer, ListBuffer} |
| import scala.collection.{Seq, mutable} |
| |
| object QueryResult { |
| def fromVertices(graph: S2Graph, vertices: Seq[S2Vertex], queryParams: Seq[QueryParam]): StepResult = { |
| val edgeWithScores = vertices.flatMap { vertex => |
| queryParams.map { queryParam => |
| val label = queryParam.label |
| val currentTs = System.currentTimeMillis() |
| val propsWithTs = Map(LabelMeta.timestamp -> |
| InnerValLikeWithTs(InnerVal.withLong(currentTs, label.schemaVersion), currentTs)) |
| |
| val edge = graph.newEdge(vertex, vertex, label, queryParam.labelWithDir.dir, propsWithTs = propsWithTs) |
| val edgeWithScore = EdgeWithScore(edge, S2Graph.DefaultScore, queryParam.label) |
| edgeWithScore |
| |
| } |
| } |
| StepResult(edgeWithScores = edgeWithScores, grouped = Nil, degreeEdges = Nil, false) |
| } |
| |
| def fromVertices(graph: S2Graph, |
| query: Query): StepResult = { |
| if (query.steps.isEmpty || query.steps.head.queryParams.isEmpty) { |
| StepResult.Empty |
| } else { |
| fromVertices(graph, query.vertices, query.steps.head.queryParams) |
| // val queryParam = query.steps.head.queryParams.head |
| // val label = queryParam.label |
| // val currentTs = System.currentTimeMillis() |
| // val propsWithTs = Map(LabelMeta.timestamp -> |
| // InnerValLikeWithTs(InnerVal.withLong(currentTs, label.schemaVersion), currentTs)) |
| // val edgeWithScores = for { |
| // vertex <- query.vertices |
| // } yield { |
| // val edge = graph.newEdge(vertex, vertex, label, queryParam.labelWithDir.dir, propsWithTs = propsWithTs) |
| // val edgeWithScore = EdgeWithScore(edge, S2Graph.DefaultScore, queryParam.label) |
| // edgeWithScore |
| // } |
| // StepResult(edgeWithScores = edgeWithScores, grouped = Nil, degreeEdges = Nil, false) |
| } |
| } |
| } |
| |
| case class QueryRequest(query: Query, |
| stepIdx: Int, |
| vertex: S2Vertex, |
| queryParam: QueryParam, |
| prevStepScore: Double = 1.0, |
| labelWeight: Double = 1.0) { |
| val nextStepOpt = |
| if (stepIdx < query.steps.size - 1) Option(query.steps(stepIdx + 1)) |
| else None |
| } |
| |
| trait WithScore[T] { |
| def score(me: T): Double |
| def withNewScore(me: T, newScore: Double): T |
| } |
| |
| object WithScore { |
| implicit val impEdgeWithScore = new WithScore[EdgeWithScore] { |
| override def score(me: EdgeWithScore): Double = me.score |
| |
| override def withNewScore(me: EdgeWithScore, newScore: Double): EdgeWithScore = me.copy(score = newScore) |
| } |
| } |
| |
| case class EdgeWithScore(edge: S2Edge, |
| score: Double, |
| label: Label, |
| orderByValues: (Any, Any, Any, Any) = StepResult.EmptyOrderByValues, |
| groupByValues: Seq[Option[Any]] = Nil, |
| stepGroupByValues: Seq[Option[Any]] = Nil, |
| filterOutValues: Seq[Option[Any]] = Nil, |
| accumulatedScores: Map[String, Double] = Map.empty) { |
| |
| def toValue(keyName: String): Option[Any] = keyName match { |
| case "score" => Option(score) |
| case _ => edge.propertyValue(keyName).map(_.innerVal.value) |
| } |
| |
| def toValues(keyNames: Seq[String]): Seq[Option[Any]] = for { |
| keyName <- keyNames |
| } yield toValue(keyName) |
| |
| } |
| |
| /** result */ |
| case class StepResult(edgeWithScores: Seq[EdgeWithScore], |
| grouped: Seq[(StepResult.GroupByKey, (Double, StepResult.Values))], |
| degreeEdges: Seq[EdgeWithScore], |
| isFailure: Boolean = false, |
| accumulatedCursors: Seq[Seq[Array[Byte]]] = Nil, |
| cursors: Seq[Array[Byte]] = Nil, |
| failCount: Int = 0) { |
| // val isInnerEmpty = innerResults.isEmpty |
| val isEmpty = edgeWithScores.isEmpty |
| } |
| |
| object StepResult { |
| |
| type Values = Seq[EdgeWithScore] |
| type GroupByKey = Seq[Option[Any]] |
| val EmptyOrderByValues = (None, None, None, None) |
| val Empty = StepResult(Nil, Nil, Nil) |
| val Failure = StepResult(Nil, Nil, Nil, true, failCount = 1) |
| |
| def mergeOrdered(left: StepResult.Values, |
| right: StepResult.Values, |
| queryOption: QueryOption): (Double, StepResult.Values) = { |
| val merged = (left ++ right) |
| val scoreSum = merged.foldLeft(0.0) { case (prev, current) => prev + current.score } |
| if (scoreSum < queryOption.scoreThreshold) (0.0, Nil) |
| else { |
| val ordered = orderBy(queryOption, merged) |
| val filtered = ordered.take(queryOption.groupBy.limit) |
| val newScoreSum = filtered.foldLeft(0.0) { case (prev, current) => prev + current.score } |
| (newScoreSum, filtered) |
| } |
| } |
| |
| def filterOutStepGroupBy(edgesWithScores: Seq[EdgeWithScore], |
| groupBy: GroupBy): Seq[EdgeWithScore] = |
| if (groupBy == GroupBy.Empty) edgesWithScores |
| else { |
| groupBy.minShouldMatch match { |
| case None => edgesWithScores |
| case Some(minShouldMatch) => |
| val MinShouldMatchParam(propKey, count, terms) = minShouldMatch |
| |
| val grouped = edgesWithScores.groupBy { edgeWithScore => |
| edgeWithScore.stepGroupByValues |
| }.filter { case (key, edges) => |
| val filtered = edges.toStream.filter{ e => |
| e.toValue(propKey) match { |
| case None => false |
| case Some(v) => terms.contains(v) |
| } |
| }.take(count) |
| |
| filtered.lengthCompare(count) >= 0 |
| } |
| |
| grouped.values.flatten.toSeq |
| } |
| } |
| |
| def orderBy(queryOption: QueryOption, notOrdered: Values): Values = { |
| import OrderingUtil._ |
| if (queryOption.withScore) { |
| val ascendingVals = if (queryOption.ascendingVals.isEmpty) QueryOption.DefaultAscendingVals else queryOption.ascendingVals |
| notOrdered.sortBy(_.orderByValues)(TupleMultiOrdering[Any](ascendingVals)) |
| } else { |
| notOrdered |
| } |
| } |
| |
| def updateScoreOnOrderByValues(scoreFieldIndex: Int, |
| orderByValues: (Any, Any, Any, Any), |
| newScore: Double): (Any, Any, Any, Any) = { |
| scoreFieldIndex match { |
| case 0 => (newScore, orderByValues._2, orderByValues._3, orderByValues._4) |
| case 1 => (orderByValues._1, newScore, orderByValues._3, orderByValues._4) |
| case 2 => (orderByValues._1, orderByValues._2, newScore, orderByValues._4) |
| case 3 => (orderByValues._1, orderByValues._2, orderByValues._3, newScore) |
| case _ => orderByValues |
| } |
| |
| } |
| |
| def toTuple4(values: Seq[Option[Any]]): (Any, Any, Any, Any) = { |
| values.length match { |
| case 1 => (values(0).getOrElse(None), None, None, None) |
| case 2 => (values(0).getOrElse(None), values(1).getOrElse(None), None, None) |
| case 3 => (values(0).getOrElse(None), values(1).getOrElse(None), values(2).getOrElse(None), None) |
| case _ => (values(0).getOrElse(None), values(1).getOrElse(None), values(2).getOrElse(None), values(3).getOrElse(None)) |
| } |
| } |
| |
| /** |
| * merge multiple StepResult into one StepResult. |
| * @param globalQueryOption |
| * @param multiStepResults |
| * @param weights |
| * @param filterOutStepResult |
| * @return |
| */ |
| def merges(globalQueryOption: QueryOption, |
| multiStepResults: Seq[StepResult], |
| weights: Seq[Double] = Nil, |
| filterOutStepResult: StepResult): StepResult = { |
| val degrees = multiStepResults.flatMap(_.degreeEdges) |
| val ls = new mutable.ListBuffer[EdgeWithScore]() |
| val agg= new mutable.HashMap[GroupByKey, ListBuffer[EdgeWithScore]]() |
| val sums = new mutable.HashMap[GroupByKey, Double]() |
| |
| |
| val filterOutSet = filterOutStepResult.edgeWithScores.foldLeft(Set.empty[Seq[Option[Any]]]) { case (prev, t) => |
| prev + t.filterOutValues |
| } |
| |
| for { |
| (weight, eachStepResult) <- weights.zip(multiStepResults) |
| (ordered, grouped) = (eachStepResult.edgeWithScores, eachStepResult.grouped) |
| } { |
| ordered.foreach { t => |
| val filterOutKey = t.filterOutValues |
| if (!filterOutSet.contains(filterOutKey)) { |
| val newScore = t.score * weight |
| val newT = t.copy(score = newScore) |
| |
| // val newOrderByValues = updateScoreOnOrderByValues(globalQueryOption.scoreFieldIdx, t.orderByValues, newScore) |
| val newOrderByValues = |
| if (globalQueryOption.orderByKeys.isEmpty) (newScore, t.edge.tsInnerVal, None, None) |
| else toTuple4(newT.toValues(globalQueryOption.orderByKeys)) |
| |
| val newGroupByValues = newT.toValues(globalQueryOption.groupBy.keys) |
| |
| ls += t.copy(score = newScore, orderByValues = newOrderByValues, groupByValues = newGroupByValues) |
| } |
| } |
| |
| // process each query's stepResult's grouped |
| for { |
| (groupByKey, (scoreSum, values)) <- grouped |
| } { |
| val buffer = agg.getOrElseUpdate(groupByKey, ListBuffer.empty[EdgeWithScore]) |
| var scoreSum = 0.0 |
| var isEmpty = true |
| values.foreach { t => |
| val filterOutKey = t.filterOutValues |
| if (!filterOutSet.contains(filterOutKey)) { |
| isEmpty = false |
| val newScore = t.score * weight |
| val newT = t.copy(score = newScore) |
| // val newOrderByValues = updateScoreOnOrderByValues(globalQueryOption.scoreFieldIdx, t.orderByValues, newScore) |
| |
| val newOrderByValues = |
| if (globalQueryOption.orderByKeys.isEmpty) (newScore, t.edge.tsInnerVal, None, None) |
| else toTuple4(newT.toValues(globalQueryOption.orderByKeys)) |
| |
| val newGroupByValues = newT.toValues(globalQueryOption.groupBy.keys) |
| |
| buffer += t.copy(score = newScore, orderByValues = newOrderByValues, groupByValues = newGroupByValues) |
| scoreSum += newScore |
| } |
| } |
| if (!isEmpty) sums += (groupByKey -> scoreSum) |
| } |
| } |
| |
| // process global groupBy |
| val (ordered, grouped) = if (globalQueryOption.groupBy.keys.nonEmpty) { |
| for { |
| edgeWithScore <- ls |
| groupByKey = edgeWithScore.groupByValues |
| } { |
| val buffer = agg.getOrElseUpdate(groupByKey, ListBuffer.empty[EdgeWithScore]) |
| buffer += edgeWithScore |
| val newScore = sums.getOrElse(groupByKey, 0.0) + edgeWithScore.score |
| sums += (groupByKey -> newScore) |
| } |
| val grouped = for { |
| (groupByKey, scoreSum) <- sums.toSeq.sortBy(_._2 * -1) |
| aggregated = agg(groupByKey) if aggregated.nonEmpty |
| sorted = orderBy(globalQueryOption, aggregated) |
| } yield (groupByKey, (scoreSum, sorted)) |
| (Nil, grouped) |
| } else { |
| val ordered = orderBy(globalQueryOption, ls) |
| (ordered, Nil) |
| } |
| |
| StepResult(edgeWithScores = ordered, grouped = grouped, degrees, failCount = multiStepResults.map(_.failCount).sum) |
| } |
| |
| //TODO: Optimize this. |
| def filterOut(graph: S2Graph, |
| queryOption: QueryOption, |
| baseStepResult: StepResult, |
| filterOutStepResult: StepResult): StepResult = { |
| |
| val filterOutSet = filterOutStepResult.edgeWithScores.foldLeft(Set.empty[Seq[Option[Any]]]) { case (prev, t) => |
| prev + t.filterOutValues |
| } |
| |
| val filteredResults = baseStepResult.edgeWithScores.filter { t => |
| val filterOutKey = t.filterOutValues |
| !filterOutSet.contains(filterOutKey) |
| } |
| |
| val grouped = for { |
| (key, (scoreSum, values)) <- baseStepResult.grouped |
| (out, in) = values.partition(v => filterOutSet.contains(v.filterOutValues)) |
| newScoreSum = scoreSum - out.foldLeft(0.0) { case (prev, current) => prev + current.score } if in.nonEmpty |
| } yield (key, (newScoreSum, in)) |
| |
| StepResult(edgeWithScores = filteredResults, grouped = grouped, baseStepResult.degreeEdges, cursors = baseStepResult.cursors, failCount = baseStepResult.failCount + filterOutStepResult.failCount) |
| } |
| } |
