1 /***********************************************************************
2  * Copyright (c) 2013-2025 General Atomics Integrated Intelligence, 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  * https://www.apache.org/licenses/LICENSE-2.0
7  ***********************************************************************/
8 
9 package org.locationtech.geomesa.arrow.io.reader
10 
11 import org.geotools.api.feature.simple.SimpleFeatureType
12 import org.geotools.api.filter.Filter
13 import org.locationtech.geomesa.arrow.features.ArrowSimpleFeature
14 import org.locationtech.geomesa.arrow.io.SimpleFeatureArrowFileReader
15 import org.locationtech.geomesa.arrow.vector.{ArrowDictionary, SimpleFeatureVector}
16 import org.locationtech.geomesa.utils.collection.CloseableIterator
17 import org.locationtech.geomesa.utils.io.CloseWithLogging
18 
19 import java.io.InputStream
20 import java.util.concurrent.atomic.AtomicBoolean
21 
22 /**
23  * Streaming reader that will re-read the input stream on each call to features()
24  *
25  * @param is input stream function
26  */
27 class MultiStreamSimpleFeatureArrowFileReader(is: () => InputStream) extends SimpleFeatureArrowFileReader {
28 
29   private val reader = new StreamingSimpleFeatureArrowFileReader(is())
30   private val createNewReader = new AtomicBoolean(false)
31 
32   override def sft: SimpleFeatureType = reader.sft
33   override def dictionaries: Map[String, ArrowDictionary] = reader.dictionaries
34   override def vectors: Seq[SimpleFeatureVector] = reader.vectors
35   override def features(filter: Filter): CloseableIterator[ArrowSimpleFeature] = {
36     // we can use the original reader for the first query, subsequent queries we have to create a new one
37     val newReader =
38       if (createNewReader.compareAndSet(false, true)) { None } else {
39         Some(new StreamingSimpleFeatureArrowFileReader(is()))
40       }
41     val iter = newReader.getOrElse(reader).features(filter)
42     new CloseableIterator[ArrowSimpleFeature] {
43       override def hasNext: Boolean = iter.hasNext
44       override def next(): ArrowSimpleFeature = iter.next()
45       override def close(): Unit = CloseWithLogging(Seq(iter) ++ newReader)
46     }
47   }
48 
49   override def close(): Unit = reader.close()
50 }