diff --git a/dataframe-excel/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/xlsx.kt b/dataframe-excel/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/xlsx.kt index c32dcffbd4..397a942336 100644 --- a/dataframe-excel/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/xlsx.kt +++ b/dataframe-excel/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/xlsx.kt @@ -306,8 +306,9 @@ public fun DataFrame.writeExcel( sheetName: String? = null, writeHeader: Boolean = true, workBookType: WorkBookType = WorkBookType.XLSX, + keepFile: Boolean = false, ) { - return writeExcel(File(path), columnsSelector, sheetName, writeHeader, workBookType) + return writeExcel(File(path), columnsSelector, sheetName, writeHeader, workBookType, keepFile) } public enum class WorkBookType { @@ -320,16 +321,22 @@ public fun DataFrame.writeExcel( sheetName: String? = null, writeHeader: Boolean = true, workBookType: WorkBookType = WorkBookType.XLSX, + keepFile: Boolean = false, ) { - val factory = when (workBookType) { - WorkBookType.XLS -> { - { HSSFWorkbook() } - } - WorkBookType.XLSX -> { - { XSSFWorkbook() } + val factory = + if (keepFile){ + when (workBookType) { + WorkBookType.XLS -> HSSFWorkbook(file.inputStream()) + WorkBookType.XLSX -> XSSFWorkbook(file.inputStream()) + } + } + else { + when (workBookType) { + WorkBookType.XLS -> HSSFWorkbook() + WorkBookType.XLSX -> XSSFWorkbook() + } } - } return file.outputStream().use { writeExcel(it, columnsSelector, sheetName, writeHeader, factory) } @@ -340,9 +347,9 @@ public fun DataFrame.writeExcel( columnsSelector: ColumnsSelector = { all() }, sheetName: String? = null, writeHeader: Boolean = true, - factory: () -> Workbook, + factory: Workbook ) { - val wb: Workbook = factory() + val wb: Workbook = factory writeExcel(wb, columnsSelector, sheetName, writeHeader) wb.write(outputStream) wb.close() diff --git a/dataframe-excel/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/XlsxTest.kt b/dataframe-excel/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/XlsxTest.kt index 976fd80e03..e6a05a087a 100644 --- a/dataframe-excel/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/XlsxTest.kt +++ b/dataframe-excel/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/XlsxTest.kt @@ -133,6 +133,21 @@ class XlsxTest { } } + @Test + fun `write to new sheet when keepFile is true`() { + val names = (1..5).map { "column$it" } + val df = dataFrameOf(names).randomDouble(7) + val fileLoc = Files.createTempFile("generated_wb", ".xlsx").toFile() + + df.writeExcel(fileLoc, sheetName = "TestSheet1") + df.writeExcel(fileLoc, sheetName = "TestSheet2", keepFile = true) + + val testSheet1Df = DataFrame.readExcel(fileLoc, sheetName = "TestSheet1") + val testSheet2Df = DataFrame.readExcel(fileLoc, sheetName = "TestSheet2") + + testSheet1Df.columnNames() shouldBe testSheet2Df.columnNames() + } + @Test fun `read xlsx file with duplicated columns and repair column names`() { shouldThrow { diff --git a/docs/StardustDocs/topics/write.md b/docs/StardustDocs/topics/write.md index ddb33fbec6..a3c590fea1 100644 --- a/docs/StardustDocs/topics/write.md +++ b/docs/StardustDocs/topics/write.md @@ -117,6 +117,21 @@ wb.close() +Add new sheets without using Apache POI directly by using a parameter to keep using the same file if it already exists + + + +```kotlin +// Create a new Excel workbook with a single sheet called "allPersons", replacing the file if it already exists -> Current sheets: allPersons +df.writeExcel(file, sheetName = "allPersons") +// Add a new sheet to the previous file without replacing it, by setting keepFile = true -> Current sheets: allPersons, happyPersons +df.filter { person -> person.isHappy }.remove("isHappy").writeExcel(file, sheetName = "happyPersons", keepFile = true) +// Add a new sheet to the previous file without replacing it, by setting keepFile = true -> Current sheets: allPersons, happyPersons, unhappyPersons +df.filter { person -> !person.isHappy }.remove("isHappy").writeExcel(file, sheetName = "unhappyPersons", keepFile = true) +``` + + + ### Writing to Apache Arrow formats Add dependency: diff --git a/tests/src/test/kotlin/org/jetbrains/kotlinx/dataframe/samples/api/Write.kt b/tests/src/test/kotlin/org/jetbrains/kotlinx/dataframe/samples/api/Write.kt index e0ce7769cf..0157f4d182 100644 --- a/tests/src/test/kotlin/org/jetbrains/kotlinx/dataframe/samples/api/Write.kt +++ b/tests/src/test/kotlin/org/jetbrains/kotlinx/dataframe/samples/api/Write.kt @@ -219,6 +219,20 @@ class Write : TestBase() { } } + @Test + fun writeXlsWithMultipleSheets() { + useTempFile { file -> + // SampleStart + // Create a new Excel workbook with a single sheet called "allPersons", replacing the file if it already exists -> Current sheets: allPersons + df.writeExcel(file, sheetName = "allPersons") + // Add a new sheet to the previous file without replacing it, by setting keepFile = true -> Current sheets: allPersons, happyPersons + df.filter { person -> person.isHappy }.remove("isHappy").writeExcel(file, sheetName = "happyPersons", keepFile = true) + // Add a new sheet to the previous file without replacing it, by setting keepFile = true -> Current sheets: allPersons, happyPersons, unhappyPersons + df.filter { person -> !person.isHappy }.remove("isHappy").writeExcel(file, sheetName = "unhappyPersons", keepFile = true) + // SampleEnd + } + } + companion object { private fun String.rejoinWithSystemLineSeparator() = rejoinWithLineSeparator(System.lineSeparator())