1 /***********************************************************************
2  * Copyright (c) 2013-2024 Commonwealth Computer Research, Inc.
3  * All rights reserved. This program and the accompanying materials
4  * are made available under the terms of the Apache License, Version 2.0
5  * which accompanies this distribution and is available at
6  * http://www.opensource.org/licenses/apache2.0.php.
7  ***********************************************************************/
8 
9 package org.locationtech.geomesa.accumulo.iterators
10 
11 import org.apache.accumulo.core.data.{Range => aRange, _}
12 import org.apache.accumulo.core.iterators.{IteratorEnvironment, SortedKeyValueIterator}
13 import org.locationtech.geomesa.accumulo.iterators.BaseAggregatingIterator.BatchScanCallback
14 import org.locationtech.geomesa.index.iterators.AggregatingScan
15 import org.locationtech.geomesa.index.iterators.AggregatingScan.{AggregateCallback, RowValue}
16 
17 /**
18  * Aggregating iterator - only works on kryo-encoded features
19  */
20 abstract class BaseAggregatingIterator[T <: AggregatingScan.Result]
21     extends SortedKeyValueIterator[Key, Value] with AggregatingScan[T] {
22 
23   import scala.collection.JavaConverters._
24 
25   var source: SortedKeyValueIterator[Key, Value] = _
26 
27   protected var topKey: Key = _
28   private var topValue: Value = new Value()
29   private var currentRange: aRange = _
30   private var needToAdvance = false
31 
32   override def init(src: SortedKeyValueIterator[Key, Value],
33                     options: java.util.Map[String, String],
34                     env: IteratorEnvironment): Unit = {
35     this.source = src
36     super.init(options.asScala.toMap)
37   }
38 
39   override def hasTop: Boolean = topKey != null
40   override def getTopKey: Key = topKey
41   override def getTopValue: Value = topValue
42 
43   override def seek(range: aRange, columnFamilies: java.util.Collection[ByteSequence], inclusive: Boolean): Unit = {
44     currentRange = range
45     source.seek(range, columnFamilies, inclusive)
46     needToAdvance = false
47     findTop()
48   }
49 
50   override def next(): Unit = findTop()
51 
52   private def findTop(): Unit = {
53     val result = aggregate(new BatchScanCallback()).result
54     if (result == null) {
55       topKey = null // hasTop will be false
56       topValue = null
57     } else {
58       if (topValue == null) {
59         // only re-create topValue if it was nulled out
60         topValue = new Value()
61       }
62       topValue.set(result)
63     }
64   }
65 
66   override protected def hasNextData: Boolean = {
67     if (needToAdvance) {
68       source.next() // advance the source iterator, this may invalidate the top key/value we've already read
69       needToAdvance = false
70     }
71     source.hasTop && !currentRange.afterEndKey(source.getTopKey)
72   }
73 
74   override protected def nextData(): RowValue = {
75     needToAdvance = true
76     topKey = source.getTopKey
77     val value = source.getTopValue.get()
78     RowValue(topKey.getRow.getBytes, 0, topKey.getRow.getLength, value, 0, value.length)
79   }
80 
81   override def deepCopy(env: IteratorEnvironment): SortedKeyValueIterator[Key, Value] =
82     throw new NotImplementedError()
83 }
84 
85 object BaseAggregatingIterator {
86 
87   private class BatchScanCallback extends AggregateCallback {
88 
89     private var bytes: Array[Byte] = _
90 
91     override def batch(bytes: Array[Byte]): Boolean = {
92       this.bytes = bytes
93       false // we want to stop scanning and return the batch
94     }
95 
96     // we always keep scanning and rely on client connections to stop the scan
97     override def partial(bytes: => Array[Byte]): Boolean = true
98 
99     def result: Array[Byte] = bytes
100   }
101 }