SQLiteTest.kt

import org.junit.Test
import java.sql.Connection
import java.sql.DriverManager
import java.sql.PreparedStatement
import java.sql.ResultSet
import kotlin.time.ExperimentalTime
import kotlin.time.measureTime

/**
 * A bunch of tests to try to understand why selecting 999 keys in a single statement is faster than selecting 1 key 999 times
 *
 * Using org.xerial:sqlite-jdbc:3.30.1
 * Run on a 2020 MacBookPro
 *
 * readChunked(                   1,         not-prepared,       no-transaction,         in) = 46.3ms
 * readChunked(                   1,         not-prepared,          transaction,         in) = 25.6ms
 * readChunked(                   1,             prepared,       no-transaction,         in) = 7.34ms
 * readChunked(                   1,             prepared,          transaction,         in) = 7.31ms
 * readChunked(                 999,         not-prepared,       no-transaction,         in) = 1.21ms
 * readChunked(                 999,         not-prepared,          transaction,         in) = 1.40ms
 * readChunked(                 999,             prepared,       no-transaction,         in) = 1.34ms
 * readChunked(                 999,             prepared,          transaction,         in) = 1.46ms
 *
 */
class MainTest {
    val total = 100_000
    val queried = 999

    @Test
    fun test() = withConnection {
        for (chunkSize in listOf(1, 999)) {
            for (prepared in listOf(false, true)) {
                for (transaction in listOf(false, true)) {
                    for (useOr in listOf(false)) {
                        readChunked(chunkSize, prepared, transaction, useOr)
                    }
                }
            }
        }
    }

    @OptIn(ExperimentalTime::class)
    private fun Connection.readChunked(chunkSize: Int, prepared: Boolean, transaction: Boolean, useOr: Boolean) {
        val inClause = "key IN (" + 0.until(chunkSize).map { "?" }.joinToString(",") + ")"
        val orClause = 0.until(chunkSize).map { "key = ?" }.joinToString(" OR ")
        val clause = if (useOr) orClause else inClause
        val sql = "SELECT key, record FROM records WHERE $clause"

        val time = measureTime {
            if (transaction) {
                execute("BEGIN TRANSACTION;")
            }
            var statement: PreparedStatement? = null
            if (prepared) {
                statement = this.prepareStatement(sql)
            }
            0.until(queried).chunked(chunkSize).forEach {
                if (!prepared) {
                    statement = this.prepareStatement(sql)
                } else {
                    statement!!.clearParameters()
                }
                0.until(chunkSize).forEach { index ->
                    statement!!.setString(index + 1, it[index].toString())
                }
                statement!!.executeQuery().apply {
                    //dump(" ")
                }.close()
            }
            if (transaction) {
                execute("COMMIT;")
            }
        }

        val preparedString = if (prepared) "prepared" else "not-prepared"
        val transactionString = if (transaction) "transaction" else "no-transaction"
        val clauseString = if (useOr) "or" else "in"
        println(
            String.format(
                "readChunked(%20s, %20s, %20s, %10s) = $time",
                chunkSize.toString(),
                preparedString,
                transactionString,
                clauseString
            )
        )
    }

    private fun Connection.execute(sql: String) {
        prepareStatement(sql).execute()
    }

    private fun withConnection(block: Connection.() -> Unit) {
        val connection = DriverManager.getConnection("jdbc:sqlite::memory:")

        connection.execute(
            """
            CREATE TABLE records (
              _id INTEGER PRIMARY KEY AUTOINCREMENT,
              key TEXT NOT NULL,
              record TEXT NOT NULL
            );
        """.trimIndent()
        )

        connection.execute("CREATE INDEX idx_records_key ON records(key);")

        repeat(total) {
            connection.execute("INSERT INTO records (key, record) VALUES ('$it', 'lorem $it ipsum lorem ipsum lorem ipsum lorem ipsum lorem ipsum');")
        }
        connection.block()
    }

    private fun ResultSet.dump(prefix: String) {
        while (next()) {
            print(prefix)
            for (i in 0.until(metaData.columnCount)) {
                print(getString(i + 1))
                print("  ")
            }
            println("")
        }
    }
}