override fun groupBy()

in src/org/jetbrains/r/visualization/inlays/dataframe/DataFrameImpl.kt [66:171]
77 lines of code
12 McCabe index (conditional complexity)

    override fun groupBy(columnName: String): DataFrame {

        val groupingColumn = columns.find { column -> column.name == columnName }!!

        // Grouping will not work currently on any Array-type column
        if(groupingColumn.type.isArray()) {
            throw Exception("Currently cannot group on array-type columns")
        }

        val alreadyProcessed = HashMap<Any?, Int>()

        val indexes = IntArray(groupingColumn.size)

        for ((i, value) in groupingColumn.withIndex()) {

            var pos = alreadyProcessed[value]

            if(pos == null) {
                pos = alreadyProcessed.size
                alreadyProcessed[value] = pos
            }

            indexes[i] = pos
        }

        val newColumnsSize = alreadyProcessed.size


        val newColumns = ArrayList<Column<*>>()

        for(column in columns) {

            if(column.name == columnName) {

                // Grouping column preserves their type

              val newColumn = when(column.type) {
                    is IntType -> {
                        val newData = IntArray(newColumnsSize)
                        alreadyProcessed.forEach{(key, value) -> newData[value] = key as Int}
                        IntColumn(column.name, newData.toCollection(ArrayList()))
                    }

                    is DoubleType -> {
                        val newData = DoubleArray(newColumnsSize)
                        alreadyProcessed.forEach{(key, value) -> newData[value] = key as Double}
                        DoubleColumn(column.name, newData.toCollection(ArrayList()))
                    }

                    is StringType -> {
                        val newData = Array(newColumnsSize) {""}
                        alreadyProcessed.forEach{(key, value) -> newData[value] = key as String}
                        StringColumn(column.name, newData.toCollection(ArrayList()))
                    }
                    else -> throw Exception("Unsupported column type ${column.type} in grouping.")
                }

                newColumns.add(newColumn)

            } else {

                val newColumn = when(column.type) {
                    is IntType -> {
                        val newData = ArrayList<ArrayList<Int>>(newColumnsSize)
                        for(i in 0 until newColumnsSize) {
                            newData.add(ArrayList())
                        }

                        for ((i, value) in (column as IntColumn).withIndex()) {
                            newData[indexes[i]].add(value)
                        }
                        IntArrayColumn(column.name, newData)
                    }

                    is DoubleType -> {
                        val newData = ArrayList<ArrayList<Double>>(newColumnsSize)
                        for(i in 0 until newColumnsSize) {
                            newData.add(ArrayList())
                        }

                        for ((i, value) in (column as DoubleColumn).withIndex()) {
                            newData[indexes[i]].add(value)
                        }
                        DoubleArrayColumn(column.name, newData)
                    }

                    is StringType -> {
                        val newData = ArrayList<ArrayList<String?>>(newColumnsSize)
                        for(i in 0 until newColumnsSize) {
                            newData.add(ArrayList())
                        }

                        for ((i, value) in (column as StringColumn).withIndex()) {
                            newData[indexes[i]].add(value)
                        }
                        StringArrayColumn(column.name, newData)
                    }
                    else -> throw Exception("Unsupported column type ${column.type} in grouping.")
                }

                newColumns.add(newColumn)
            }
        }

        return DataFrameImpl(newColumns)
    }