Add MongoDB support

[AndreiKingsley]

[DestBro]

Since MongoDB is used sometimes for BigData things, a fast conversion of BSON Documents into a DataFrame should be a priority. Right now I do it like this:

@file:DependsOn("org.mongodb:mongodb-driver-kotlin-sync:5.5.1")

import com.mongodb.kotlin.client.MongoClient
import com.mongodb.kotlin.client.MongoDatabase

val connectionString = "mongodb://root:example@localhost:27017"
val client = MongoClient.create(connectionString = connectionString)

val databaseName = "db name"
val db: MongoDatabase = client.getDatabase(databaseName = databaseName)
var collection: MongoCollection<Document> = db.getCollection("comparison_result")

val dataframe = collection.find().allowDiskUse(true).map { DataFrame.readJsonStr(it.toJson()).flatten()  }.toList().toDataFrame().explode("value").flatten("value")

The current approach took about 4 minutes to load approximately 32,000 records (with nested arrays). I'm not sure excatly where is the bottleneck (network IO, json de/serialisation, etc) - because it was just an experiment.

Some comments about this solution from @AndreiKingsley :

As for your current approach — nice workaround! One small improvement you might consider: instead of fully deserializing and constructing manually, you can use DataRow.readJsonStr() on each row, and then combine everything using List.concat() to form the final DataFrame. That should simplify the code a bit and possibly improve performance too.

However I think that conversion of BSON to JSON format should not be part of the solution, maybe there is another way.

[Jolanrensen]

DataFrame works very well when converting a list of data classes to a dataframe. I see there are many ways in MongoDB to instantiate Kotlin classes from the database: data classes, kotlinx serialization. If you get an iterable or list of classes, you can just call .toDataFrame(). It's also possible to combine our gradle/ksp plugin or compiler plugin by annotating the data classes with @DataSchema.

Finally, I see it's possible to create your own Codec for MongoDB. I haven't explored this yet, but maybe it might be possible to go directly from MongoDB to DataFrame using a custom codec. Note that this may be impossible if they want you to use their own collection type and you're only allowed to specify the types inside. In that case, the data class approach will be best :)

[DestBro]

That's true I can explore more ways to optimize it, but this particular issue in github was created with an intention of having official MongoDB support, something similar:

But not just for SQL Databases, but also MongoDB, maybe create a separate MongoDataFrame (inspired from GeoDataFrame), of course it makes only sense if you agree to support MongoDB in first place.

[DestBro]

I will check when I have time what solution is the fastest.

[Jolanrensen]

@DestBro True :) but it would be nice to have a proof of concept first

[DestBro]

@Jolanrensen
I continue to experiment, the mongo Document also has the toMap() method.
Since there is no way to create a DataFrame from a List of Maps, I got inspired from this solution: #90, and made something like:

import org.bson.Document
import com.mongodb.kotlin.client.coroutine.MongoCollection
import kotlinx.coroutines.flow.toList
import org.jetbrains.kotlinx.dataframe.api.*
import org.jetbrains.kotlinx.dataframe.DataFrame
import kotlin.reflect.KProperty

suspend fun MongoCollection<Document>.toDataFrame(): DataFrame<Any?> {
    val docs: List<Map<String, Any>> = this.find().map { it.toMap() }.toList()
    return docs.toDataFrame()
}

fun List<Map<String, Any?>>.toDataFrame(): AnyFrame {
    val columns = mutableMapOf<String, MutableList<Any?>>()
    val notNullCols = mutableSetOf<String>()
    val columnSize = size

    forEachIndexed { rowIndex, row ->
        for (col in row.keys) {
            if (columns[col] == null)
                columns[col] = mutableListOf()

            val value =
                if (row[col].let { it is String && it.isEmpty() }) null
                else row[col]

            if (value != null)
                notNullCols += col

            while (columns[col]!!.size < rowIndex)
                columns[col]!! += null

            columns[col]!! += value
        }
    }

    return columns
        .filter { it.value.isNotEmpty() && it.key in notNullCols }
        .also { map ->
            for ((_, value) in map) {
                while (value.size < columnSize)
                    value += null
            }
        }
        .map { (key, value) -> DataColumn.createByType(key, value, infer = Infer.Type) }
        .toDataFrame()
}

It took 2 seconds to parse the Documents into a DataFrame, and the schema is very similar (even better) to the Json solution above.

To mention again, there are serious performance issue with: [Json] -> [DataRows] -> DataFrame

Also I checked the solution from this issue #719, it got wrong schema and values for my use-case.

[DestBro]

I see it's possible to create your own Codec for MongoDB. I haven't explored this yet, but maybe it might be possible to go directly from MongoDB to DataFrame using a custom codec.

if it's possible to use DataFrame instead of Documents directly using Codecs, I think it will give more control and avoid conversions on the way: Document -> [Map] --(unofficial convertor)-> DataFrame, and maybe avoiding the guessing/Inferring the type too: DataColumn.createByType(key, value, infer = Infer.Type)

[DestBro]

#1309 (comment) The solution had flaws, it did not handle nested Documents, here is another impl:

fun recursiveToDataFrame(documents: List<Map<String, Any?>>): DataFrame<*> {
    val allKeys = documents.flatMap { it.keys }.toSet()
    val columnarMap = allKeys.associateWith { key ->
        documents.map { document ->
            val value = document[key]
            transformNestedValue(value)
        }
    }
    return columnarMap.toDataFrame()
}

fun transformNestedValue(value: Any?): Any? {
    return when (value) {
        is Map<*, *> -> {
            val nestedMap = value as Map<String, Any?>
            if (nestedMap.isEmpty()) {
                DataFrame.Empty
            } else {
                recursiveToDataFrame(listOf(nestedMap))
            }
        }
        is List<*> -> {
            val nestedDocuments = value.mapNotNull { item ->
                when (item) {
                    is Map<*, *> -> item as Map<String, Any?>
                    else -> null
                }
            }
            if (nestedDocuments.isNotEmpty()) {
                recursiveToDataFrame(nestedDocuments)
            } else {
                value.map { transformNestedValue(it) }
            }
        }
        else -> value
    }
}

suspend fun MongoCollection<Document>.toDataFrame() = recursiveToDataFrame(find().map { it.toMap() }.toList())

runBlocking {
    mongoCollectionComparison.toDataFrame()
}

But it still has some differences compared to the Json approach:

• Some specific Mongo types are represented differently: org.bson.types.Binary@ac5f8174 instead of "NeMAFnMAACUjQA==" (base64 of binary)