Skip to content

Commit 5e9c220

Browse files
PaulanerusPaulanerus
authored
Merge pull request #50 from Paulanerus/dev
Dev
2 parents 0626ff6 + 086f85a commit 5e9c220

19 files changed

Lines changed: 534 additions & 109 deletions

File tree

‎.github/workflows/build-and-publish.yml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,8 @@ jobs:
108108
- name: Set up Java (${{ matrix.distribution }})
109109
uses: actions/setup-java@v5.0.0
110110
with:
111-
java-version: '21'
111+
java-version: '21.0.8'
112+
check-latest: true
112113
distribution: ${{ matrix.distribution }}
113114
cache: 'gradle'
114115

‎api/src/main/kotlin/dev/paulee/api/data/IDataService.kt‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,10 @@ interface IDataService : Closeable {
1010

1111
suspend fun createDataPool(dataInfo: DataInfo, onProgress: (progress: Int) -> Unit): Boolean
1212

13+
suspend fun deleteDataPool(name: String): Boolean
14+
15+
suspend fun rebuildDataPool(dataInfo: DataInfo, onProgress: (progress: Int) -> Unit): Boolean
16+
1317
fun selectDataPool(selection: String)
1418

1519
fun getSelectedPool(): String
@@ -44,4 +48,4 @@ interface IDataService : Closeable {
4448
fun dataDir(): Path
4549

4650
fun modelDir(): Path
47-
}
51+
}

‎api/src/main/kotlin/dev/paulee/api/data/provider/StorageProvider.kt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ interface IStorageProvider : Closeable {
2626

2727
fun get(
2828
name: String,
29-
ids: List<Long> = emptyList(),
29+
ids: LinkedHashSet<Long> = LinkedHashSet(),
3030
whereClause: List<String> = emptyList(),
3131
filter: List<String> = emptyList(),
3232
order: QueryOrder? = null,
@@ -36,7 +36,7 @@ interface IStorageProvider : Closeable {
3636

3737
fun count(
3838
name: String,
39-
ids: List<Long> = emptyList(),
39+
ids: LinkedHashSet<Long> = LinkedHashSet(),
4040
whereClause: List<String> = emptyList(),
4141
filter: List<String> = emptyList(),
4242
): Long

‎core/src/main/kotlin/dev/paulee/core/data/DataServiceImpl.kt‎

Lines changed: 114 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -6,18 +6,22 @@ import dev.paulee.api.data.*
66
import dev.paulee.api.data.provider.IStorageProvider
77
import dev.paulee.api.data.provider.ProviderStatus
88
import dev.paulee.api.data.provider.QueryOrder
9+
import dev.paulee.api.data.provider.StorageType
910
import dev.paulee.api.internal.Embedding
1011
import dev.paulee.core.data.analysis.Indexer
1112
import dev.paulee.core.data.model.DataPool
1213
import dev.paulee.core.data.provider.EmbeddingProvider
1314
import dev.paulee.core.data.provider.StorageProvider
1415
import dev.paulee.core.splitStr
1516
import kotlinx.coroutines.Dispatchers
17+
import kotlinx.coroutines.sync.Mutex
18+
import kotlinx.coroutines.sync.withLock
1619
import kotlinx.coroutines.withContext
1720
import org.slf4j.LoggerFactory.getLogger
1821
import java.io.IOException
1922
import java.nio.file.Path
2023
import java.time.Duration
24+
import java.util.concurrent.ConcurrentHashMap
2125
import kotlin.coroutines.cancellation.CancellationException
2226
import kotlin.io.path.*
2327
import kotlin.math.ceil
@@ -49,7 +53,9 @@ object DataServiceImpl : IDataService {
4953

5054
private val storageProvider = mutableMapOf<String, IStorageProvider>()
5155

52-
private val dataPools = mutableMapOf<String, DataPool>()
56+
private val dataPools = ConcurrentHashMap<String, DataPool>()
57+
58+
private val mutex = Mutex()
5359

5460
init {
5561
loadDataPools(FileService.dataDir)
@@ -100,7 +106,7 @@ object DataServiceImpl : IDataService {
100106
currentProvider.streamData(name)
101107
.chunked(BATCH_SIZE)
102108
.forEach { entries ->
103-
dataPool.indexer.indexEntries(name, entries)
109+
dataPool.indexer?.indexEntries(name, entries)
104110

105111
processedBatches++
106112

@@ -114,7 +120,7 @@ object DataServiceImpl : IDataService {
114120
}
115121
}
116122

117-
dataPool.indexer.finish()
123+
dataPool.indexer?.finish()
118124
EmbeddingProvider.finish()
119125

120126
dataPools[dataInfo.name] = dataPool
@@ -133,6 +139,84 @@ object DataServiceImpl : IDataService {
133139
}
134140
}
135141

142+
@OptIn(ExperimentalPathApi::class)
143+
override suspend fun deleteDataPool(name: String): Boolean = withContext(Dispatchers.IO) {
144+
val (pool, wasCurrent, externalProvider) = mutex.withLock {
145+
val pool = dataPools.remove(name) ?: return@withLock null
146+
147+
val isCurrent = currentPool == name
148+
149+
if (isCurrent) {
150+
currentPool = null
151+
currentField = null
152+
}
153+
154+
val externalProvider = storageProvider.remove(name)
155+
156+
Triple(pool, isCurrent, externalProvider)
157+
} ?: return@withContext false
158+
159+
runCatching { pool.storageProvider?.close() }
160+
.onFailure { e -> logger.error("Failed to close storage provider for $name.", e) }
161+
162+
runCatching { pool.indexer?.close() }
163+
.onFailure { e -> logger.error("Failed to close indexer for $name.", e) }
164+
165+
externalProvider?.let {
166+
runCatching { it.close() }
167+
.onFailure { e -> logger.error("Failed to close external storage provider for $name.", e) }
168+
}
169+
170+
val poolPath = FileService.dataDir.resolve(name)
171+
172+
if (poolPath.notExists()) return@withContext true
173+
174+
runCatching { poolPath.deleteRecursively() }
175+
.onFailure { e ->
176+
logger.error("Failed to delete directory for $name.", e)
177+
return@withContext false
178+
}
179+
180+
logger.info("Deleted data pool $name.")
181+
182+
if (wasCurrent && dataPools.isNotEmpty()) {
183+
mutex.withLock {
184+
val newPool = dataPools.entries.firstOrNull() ?: return@withLock
185+
186+
currentPool = newPool.key
187+
currentField = newPool.value.defaultClass
188+
}
189+
}
190+
191+
true
192+
}
193+
194+
override suspend fun rebuildDataPool(dataInfo: DataInfo, onProgress: (progress: Int) -> Unit): Boolean {
195+
logger.info("Rebuilding data pool ${dataInfo.name}.")
196+
197+
val (oldPool, oldField) = mutex.withLock { currentPool to currentField }
198+
199+
var info = dataInfo
200+
201+
if (info.storageType == StorageType.SQLITE) {
202+
logger.warn("${info.name} uses deprecated SQLite storage type. Replacing with default storage type.")
203+
info = info.copy(storageType = StorageType.Default)
204+
}
205+
206+
if (!deleteDataPool(info.name)) return false
207+
208+
val created = createDataPool(info, onProgress)
209+
210+
if (created && oldPool == info.name) {
211+
mutex.withLock {
212+
currentPool = oldPool
213+
currentField = oldField
214+
}
215+
}
216+
217+
return created
218+
}
219+
136220
@OptIn(ExperimentalPathApi::class)
137221
fun loadDataPools(path: Path) {
138222
if (path.notExists()) {
@@ -153,15 +237,25 @@ object DataServiceImpl : IDataService {
153237
val dataInfo = FileService.fromJson(runCatching { jsonFile.readText() }.getOrDefault(""))
154238
?: return@forEachDirectoryEntry
155239

240+
val infoName = dataInfo.name
241+
242+
if (dataInfo.storageType == StorageType.SQLITE) {
243+
logger.warn("Data pool '$infoName' uses deprecated SQLite storage type. Skipping.")
244+
245+
dataPools[infoName] = DataPool(null, dataInfo, null)
246+
247+
return@forEachDirectoryEntry
248+
}
249+
156250
val storageProvider = StorageProvider.of(dataInfo.storageType)
157251

158252
if (storageProvider.init(dataInfo, child) == ProviderStatus.Exists) {
159-
dataPools[dataInfo.name] =
253+
dataPools[infoName] =
160254
DataPool(Indexer(child.resolve("index"), dataInfo), dataInfo, storageProvider)
161255

162-
logger.info("Loaded ${dataInfo.name} data pool.")
256+
logger.info("Loaded $infoName data pool.")
163257
} else {
164-
logger.info("Deleting invalid or empty data pool directory '${dataInfo.name}'.")
258+
logger.info("Deleting invalid or empty data pool directory '$infoName'.")
165259

166260
runCatching { child.deleteRecursively() }
167261
.onFailure { e -> logger.error("Failed to delete directory ${child.fileName}.", e) }
@@ -210,14 +304,15 @@ object DataServiceImpl : IDataService {
210304

211305
val dataPool = this.dataPools[this.currentPool] ?: return emptyList()
212306

307+
if (dataPool.storageProvider == null) return emptyList()
308+
213309
val fieldExists = dataPool.dataInfo.sources
214310
.firstOrNull { it.name == current }
215311
?.fields
216312
?.any { it.name == field } == true
217313

218314
if (!fieldExists) return emptyList()
219315

220-
221316
return dataPool.storageProvider.suggestions(current, field, value, 6)
222317
}
223318

@@ -236,6 +331,8 @@ object DataServiceImpl : IDataService {
236331

237332
val dataPool = this.dataPools[this.currentPool] ?: return Pair(emptyList(), emptyMap())
238333

334+
if (dataPool.storageProvider == null) return Pair(emptyList(), emptyMap())
335+
239336
val (filterQuery, filter) = this.getPreFilter(query)
240337

241338
val indexResult = dataPool.search(this.handleReplacements(dataPool.metadata, filterQuery), isSemantic)
@@ -279,6 +376,8 @@ object DataServiceImpl : IDataService {
279376

280377
val dataPool = this.dataPools[this.currentPool] ?: return Triple(-1, -1, emptySet())
281378

379+
if (dataPool.storageProvider == null) return Triple(-1, -1, emptySet())
380+
282381
val (filterQuery, filter) = this.getPreFilter(query)
283382

284383
val indexResult = dataPool.search(handleReplacements(dataPool.metadata, filterQuery), isSemantic)
@@ -298,8 +397,8 @@ object DataServiceImpl : IDataService {
298397
override fun createStorageProvider(infoName: String, path: Path): IStorageProvider? {
299398
val dataInfo = this.dataPools[infoName]?.dataInfo ?: return null
300399

301-
if (this.storageProvider[infoName] == null) this.storageProvider[infoName] =
302-
StorageProvider.of(dataInfo.storageType)
400+
if (this.storageProvider[infoName] == null)
401+
this.storageProvider[infoName] = StorageProvider.of(dataInfo.storageType)
303402

304403
val provider = this.storageProvider[infoName]
305404

@@ -327,8 +426,8 @@ object DataServiceImpl : IDataService {
327426
this.storageProvider.forEach { it.value.close() }
328427

329428
this.dataPools.values.forEach {
330-
it.storageProvider.close()
331-
it.indexer.close()
429+
it.storageProvider?.close()
430+
it.indexer?.close()
332431
}
333432

334433
EmbeddingProvider.close()
@@ -337,6 +436,8 @@ object DataServiceImpl : IDataService {
337436
private fun handleReplacements(replacements: Map<String, Any>, query: String): String {
338437
val dataPool = this.dataPools[this.currentPool] ?: return query
339438

439+
if (dataPool.storageProvider == null) return query
440+
340441
return variantPattern.replace(query) {
341442
val key = it.groupValues[1]
342443
val value = it.groupValues[2]
@@ -356,6 +457,8 @@ object DataServiceImpl : IDataService {
356457
private fun getPreFilter(query: String): Pair<String, List<String>> {
357458
val dataPool = this.dataPools[this.currentPool] ?: return Pair(query, emptyList())
358459

460+
if (dataPool.storageProvider == null) return Pair(query, emptyList())
461+
359462
val filters = preFilterPattern.findAll(query).map { it.value }.toSet()
360463

361464
val queryWithoutFilter = filters.fold(query) { acc, filter -> acc.replace(filter, "") }.trim()

‎core/src/main/kotlin/dev/paulee/core/data/FileService.kt‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,12 @@ internal object FileService {
9191
private val mapper = jacksonObjectMapper().apply { enable(SerializationFeature.INDENT_OUTPUT) }
9292

9393
init {
94-
logger.info("Operating system: ${Platform.current} | CUDA: ${Platform.isCuda12xInstalled}, cuDNN: ${Platform.isCuDNNInstalled}")
94+
val osSpecific = when (Platform.current) {
95+
Platform.MacOS -> ""
96+
else -> "| CUDA: ${Platform.isCuda12xInstalled}, cuDNN: ${Platform.isCuDNNInstalled}"
97+
}
98+
99+
logger.info("Operating system: ${Platform.current} $osSpecific")
95100
}
96101

97102
fun toJson(dataInfo: DataInfo): String? = runCatching { this.mapper.writeValueAsString(dataInfo) }

‎core/src/main/kotlin/dev/paulee/core/data/analysis/Indexer.kt‎

Lines changed: 42 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import org.apache.lucene.search.IndexSearcher
1818
import org.apache.lucene.search.Query
1919
import org.apache.lucene.store.BaseDirectory
2020
import org.apache.lucene.store.FSDirectory
21+
import org.apache.lucene.util.Version
2122
import org.slf4j.LoggerFactory.getLogger
2223
import java.io.Closeable
2324
import java.nio.file.Path
@@ -70,6 +71,8 @@ internal class Indexer(path: Path, dataInfo: DataInfo) : Closeable {
7071
private val embeddingFields = mutableMapOf<String, Embedding.Model>()
7172

7273
init {
74+
checkForVersionCompatibility()
75+
7376
dataInfo.sources.forEach {
7477
val normalized = normalizeDataSource(it.name)
7578

@@ -180,11 +183,7 @@ internal class Indexer(path: Path, dataInfo: DataInfo) : Closeable {
180183

181184
if (normalized.isBlank()) return emptyList()
182185

183-
DirectoryReader.openIfChanged(this.reader)?.let {
184-
this.reader.close()
185-
186-
this.reader = it
187-
}
186+
refreshReader()
188187

189188
val searcher = IndexSearcher(this.reader)
190189

@@ -208,11 +207,7 @@ internal class Indexer(path: Path, dataInfo: DataInfo) : Closeable {
208207
fun searchMatchingVec(field: String, query: String, similarity: Float): List<Document> {
209208
val model = embeddingFields[field] ?: return emptyList()
210209

211-
DirectoryReader.openIfChanged(this.reader)?.let {
212-
this.reader.close()
213-
214-
this.reader = it
215-
}
210+
if (query.isBlank()) return emptyList()
216211

217212
val searcher = IndexSearcher(this.reader)
218213

@@ -237,6 +232,12 @@ internal class Indexer(path: Path, dataInfo: DataInfo) : Closeable {
237232
.replace(TRAILING_REGEX, "")
238233
}
239234

235+
private fun refreshReader() =
236+
DirectoryReader.openIfChanged(reader)?.let {
237+
reader.close()
238+
reader = it
239+
}
240+
240241
private fun createDoc(
241242
name: String,
242243
map: Map<String, String>,
@@ -259,5 +260,35 @@ internal class Indexer(path: Path, dataInfo: DataInfo) : Closeable {
259260
}
260261
}
261262
}
262-
}
263263

264+
private fun checkForVersionCompatibility() {
265+
if (!DirectoryReader.indexExists(directory)) return
266+
267+
val currentMajor = Version.LATEST.major
268+
269+
val segInfo = runCatching { SegmentInfos.readLatestCommit(directory) }.getOrElse {
270+
logger.error("Exception: Failed to read SegmentInfos.", it)
271+
return
272+
}
273+
274+
val indexMajorVersion = segInfo.indexCreatedVersionMajor
275+
276+
if (currentMajor == indexMajorVersion) return
277+
278+
if (currentMajor < indexMajorVersion) {
279+
logger.warn("Indexer version ($indexMajorVersion) is newer than current version ($currentMajor).")
280+
return
281+
}
282+
283+
if (currentMajor - 1 > indexMajorVersion) {
284+
logger.error("Index is at least two major versions behind current version.")
285+
return
286+
}
287+
288+
logger.warn("Index version ($indexMajorVersion) is older than current version ($currentMajor) and will be upgraded.")
289+
290+
runCatching { IndexUpgrader(directory).upgrade() }
291+
.onSuccess { logger.info("Indexer upgraded successfully.") }
292+
.onFailure { logger.error("Failed to upgrade Indexer from $indexMajorVersion to $currentMajor.", it) }
293+
}
294+
}

0 commit comments

Comments
 (0)