Merge pull request #153 from Kotlin/fix-read-excel

use index math for skip rows #132

Author: koperagen

dataframe-excel/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/xlsx.kt

@@ -8,6 +8,7 @@ import org.apache.poi.ss.usermodel.Cell
 import org.apache.poi.ss.usermodel.CellType
 import org.apache.poi.ss.usermodel.DateUtil
 import org.apache.poi.ss.usermodel.RichTextString
+import org.apache.poi.ss.usermodel.Row
 import org.apache.poi.ss.usermodel.Sheet
 import org.apache.poi.ss.usermodel.Workbook
 import org.apache.poi.ss.usermodel.WorkbookFactory
@@ -148,7 +149,7 @@ public fun DataFrame.Companion.readExcel(
     skipRows: Int = 0,
     rowsCount: Int? = null
 ): AnyFrame {
-    val columnIndexes = if (columns != null) {
+    val columnIndexes: Iterable<Int> = if (columns != null) {
         columns.split(",").flatMap {
             if (it.contains(":")) {
                 val (start, end) = it.split(":").map { CellReference.convertColStringToIndex(it) }
@@ -158,42 +159,58 @@ public fun DataFrame.Companion.readExcel(
             }
         }
     } else {
-        sheet.getRow(skipRows).map { it.columnIndex }
+        val headerRow = checkNotNull(sheet.getRow(skipRows)) {
+            "Row number ${skipRows + 1} (1-based index) is not defined on the sheet ${sheet.sheetName}"
+        }
+        val firstCellNum = headerRow.firstCellNum
+        check(firstCellNum != (-1).toShort()) {
+            "There are no defined cells on header row number ${skipRows + 1} (1-based index). Pass `columns` argument to specify what columns to read or make sure the index is correct"
+        }
+        headerRow.firstCellNum until headerRow.lastCellNum
     }
 
-    val headerRow = sheet.getRow(skipRows)
-    val valueRows = sheet.drop(1 + skipRows).let { if (rowsCount != null) it.take(rowsCount) else it }
+    val headerRow: Row? = sheet.getRow(skipRows)
+    val first = skipRows + 1
+    val last = rowsCount?.let { first + it - 1 } ?: sheet.lastRowNum
+    val valueRowsRange = (first..last)
+
     val columns = columnIndexes.map { index ->
-        val headerCell = headerRow.getCell(index)
+        val headerCell = headerRow?.getCell(index)
         val name = if (headerCell?.cellType == CellType.NUMERIC) {
             headerCell.numericCellValue.toString() // Support numeric-named columns
         } else {
             headerCell?.stringCellValue ?: CellReference.convertNumToColString(index) // Use Excel column names if no data
         }
-        val values: List<Any?> = valueRows.map {
-            val cell: Cell? = it.getCell(index)
-            when (cell?.cellType) {
-                CellType._NONE -> error("Cell ${cell.address} of sheet ${sheet.sheetName} has a CellType that should only be used internally. This is a bug, please report https://github.com/Kotlin/dataframe/issues")
-                CellType.NUMERIC -> {
-                    val number = cell.numericCellValue
-                    when {
-                        DateUtil.isCellDateFormatted(cell) -> DateUtil.getLocalDateTime(number).toKotlinLocalDateTime()
-                        else -> number
-                    }
-                }
-                CellType.STRING -> cell.stringCellValue
-                CellType.FORMULA -> cell.numericCellValue
-                CellType.BLANK -> cell.stringCellValue
-                CellType.BOOLEAN -> cell.booleanCellValue
-                CellType.ERROR -> cell.errorCellValue
-                null -> null
-            }
+
+        val values: List<Any?> = valueRowsRange.map {
+            val row: Row? = sheet.getRow(it)
+            val cell: Cell? = row?.getCell(index)
+            cell.cellValue(sheet.sheetName)
         }
         DataColumn.createWithTypeInference(name, values)
     }
     return dataFrameOf(columns)
 }
 
+private fun Cell?.cellValue(sheetName: String): Any? =
+    when (this?.cellType) {
+        CellType._NONE -> error("Cell $address of sheet $sheetName has a CellType that should only be used internally. This is a bug, please report https://github.com/Kotlin/dataframe/issues")
+        CellType.NUMERIC -> {
+            val number = numericCellValue
+            when {
+                DateUtil.isCellDateFormatted(this) -> DateUtil.getLocalDateTime(number).toKotlinLocalDateTime()
+                else -> number
+            }
+        }
+
+        CellType.STRING -> stringCellValue
+        CellType.FORMULA -> numericCellValue
+        CellType.BLANK -> stringCellValue
+        CellType.BOOLEAN -> booleanCellValue
+        CellType.ERROR -> errorCellValue
+        null -> null
+    }
+
 public fun <T> DataFrame<T>.writeExcel(
     path: String,
     columnsSelector: ColumnsSelector<T, *> = { all() },

dataframe-excel/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/XlsxTest.kt

@@ -1,12 +1,15 @@
 package org.jetbrains.kotlinx.dataframe.io
 
+import io.kotest.assertions.throwables.shouldThrow
 import io.kotest.matchers.shouldBe
 import kotlinx.datetime.LocalDateTime
 import org.apache.poi.ss.usermodel.WorkbookFactory
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.api.concat
 import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
 import org.jetbrains.kotlinx.dataframe.api.toColumn
+import org.jetbrains.kotlinx.dataframe.impl.DataFrameSize
+import org.jetbrains.kotlinx.dataframe.size
 import org.junit.Test
 import java.net.URL
 import java.nio.file.Files
@@ -83,11 +86,27 @@ class XlsxTest {
     fun `read header on second row`() {
         val df = DataFrame.readExcel(testResource("custom_header_position.xlsx"), skipRows = 1)
         df.columnNames() shouldBe listOf("header1", "header2")
+        df.size() shouldBe DataFrameSize(2, 3)
     }
 
     @Test
     fun `consider skipRows when obtaining column indexes`() {
-        val df = DataFrame.readExcel(testResource("header.xlsx"), skipRows = 6)
+        val df = DataFrame.readExcel(testResource("header.xlsx"), skipRows = 6, rowsCount = 1)
         df.columnNames() shouldBe listOf("Well", "Well Position", "Omit", "Sample Name", "Target Name", "Task", "Reporter", "Quencher")
+        df shouldBe dataFrameOf("Well", "Well Position", "Omit", "Sample Name", "Target Name", "Task", "Reporter", "Quencher")(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0,)
+    }
+
+    @Test
+    fun `use indexes math to skip rows`() {
+        val df = DataFrame.readExcel(testResource("repro.xls"), skipRows = 4)
+        df.columnNames() shouldBe listOf("a")
+        df.rowsCount() shouldBe 2
+    }
+
+    @Test
+    fun `throw when there are no defined cells on header row`() {
+        shouldThrow<IllegalStateException> {
+            DataFrame.readExcel(testResource("xlsx6.xlsx"), skipRows = 4)
+        }
     }
 }