xtdb.vector.RelationWriter.kt Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of xtdb-core Show documentation

An open source document database with bitemporal graph queries

There is a newer version: 2.0.0-beta4

package xtdb.vector

import org.apache.arrow.memory.BufferAllocator
import org.apache.arrow.vector.types.pojo.FieldType
import xtdb.arrow.Relation
import xtdb.arrow.Vector
import xtdb.arrow.VectorPosition

@Suppress("unused")
class RelationWriter(private val allocator: BufferAllocator) : IRelationWriter {
    private val wp = VectorPosition.build()
    private val writers = mutableMapOf()

    constructor(allocator: BufferAllocator, writers: List) : this(allocator) {
        this.writers.putAll(writers.associateBy { it.vector.name })
    }

    override fun iterator() = writers.iterator()

    override fun writerPosition() = wp

    override fun startRow() = Unit

    override fun endRow() {
        val pos = ++wp.position
        writers.values.forEach { it.populateWithAbsents(pos) }
    }

    override fun colWriter(colName: String) = writers[colName] ?: colWriter(colName, UNION_FIELD_TYPE)

    override fun colWriter(colName: String, fieldType: FieldType) =
        writers[colName]
            ?.also { it.checkFieldType(fieldType) }

            ?: writerFor(fieldType.createNewSingleVector(colName, allocator, null))
                .also {
                    it.populateWithAbsents(wp.position)
                    writers[colName] = it
                }

    override fun openAsRelation() = Relation(writers.values.map { Vector.fromArrow(it.vector) }, wp.position)

    override fun close() {
        writers.values.forEach(IVectorWriter::close)
    }
}