Skip to content

Add methods to stream generated rows directly to file #2468

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
labkey-tchad merged 2 commits into from
May 30, 2025
Merged

Add methods to stream generated rows directly to file #2468

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
847f0d4
Add methods to stream generated rows directly to file
labkey-tchad May 29, 2025
9e4b110
remove out-of-scope changes
labkey-tchad May 30, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions src/org/labkey/test/tests/assay/UploadLargeExcelAssayTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ protected void doCleanup(boolean afterTest)
@BeforeClass
public static void setupProject() throws Exception
{
UploadLargeExcelAssayTest init = (UploadLargeExcelAssayTest) getCurrentTest();
UploadLargeExcelAssayTest init = getCurrentTest();

init.doSetup();
}
Expand Down Expand Up @@ -85,9 +85,8 @@ public void testUpload200kRows() throws Exception
String fileName = "200kXlsxFile.xlsx";
var dgen = new TestDataGenerator("samples", "chaos_sample", getProjectName())
.withColumns(ASSAY_FIELDS);
dgen.generateRows(200_000);
log("writing large .xlsx file");
var largeExcelFile = dgen.writeData(fileName);
var largeExcelFile = dgen.writeData(fileName, 200_000);
log("finished writing large .xlsx file");

// import large generated excel to assay1
Expand Down
170 changes: 106 additions & 64 deletions src/org/labkey/test/util/TestDataGenerator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,11 +16,8 @@
package org.labkey.test.util;

import org.apache.commons.csv.CSVFormat;
import org.apache.commons.io.FileUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.time.DateUtils;
import org.apache.poi.xssf.streaming.SXSSFRow;
import org.apache.poi.xssf.streaming.SXSSFWorkbook;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.junit.Assert;
Expand All @@ -37,15 +34,13 @@
import org.labkey.remoteapi.query.SelectRowsResponse;
import org.labkey.remoteapi.query.Sort;
import org.labkey.serverapi.reader.TabLoader;
import org.labkey.test.TestFileUtils;
import org.labkey.test.WebTestHelper;
import org.labkey.test.params.FieldDefinition;
import org.labkey.test.util.data.ColumnNameMapper;
import org.labkey.test.util.data.TestDataUtils;
import org.labkey.test.util.query.QueryApiHelper;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
Expand All @@ -54,8 +49,11 @@
import java.util.Collections;
import java.util.Date;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.ThreadLocalRandom;
import java.util.function.Supplier;
Expand Down Expand Up @@ -313,7 +311,7 @@ public TestDataGenerator withGeneratedRows(int desiredRowCount)

public TestDataGenerator addDataSupplier(String columnName, Supplier<Object> supplier)
{
_dataSuppliers.put(columnName, ()-> supplier.get());
_dataSuppliers.put(columnName, supplier);
return this;
}

Expand Down Expand Up @@ -400,7 +398,7 @@ public List<PropertyDescriptor> getColumns()

public void generateRows(int numberOfRowsToGenerate)
{
if (_columns.keySet().size() == 0)
if (_columns.isEmpty())
throw new IllegalStateException("can't generate row data without column definitions");

for (int i= 0; i < numberOfRowsToGenerate; i++)
Expand Down Expand Up @@ -444,7 +442,7 @@ private Supplier<Object> getDefaultDataSupplier(String columnType)
case "double":
return ()-> randomDouble(0, 20);
case "boolean":
return ()-> randomBoolean();
return this::randomBoolean;
case "date":
case "datetime":
return ()-> randomDateString(DateUtils.addWeeks(new Date(), -39), new Date());
Expand Down Expand Up @@ -670,40 +668,6 @@ public String getDataAsTsv()
return TestDataUtils.stringFromRowMaps(_rows, getFieldsForFile(), true, CSVFormat.TDF);
}

public File writeGeneratedDataToExcel(String sheetName, String fileName) throws IOException
{
File file = new File(TestFileUtils.getTestTempDir(), fileName);
FileUtils.forceMkdirParent(file);

try (SXSSFWorkbook workbook = new SXSSFWorkbook(1000); // only holds 1000 rows in memory
FileOutputStream out = new FileOutputStream(file))
{
var sheet = workbook.createSheet(sheetName);

// write headers as row 0
List<String> columnNames = getFieldsForFile();
var headerRow = sheet.createRow(0);
for (int i = 0; i < columnNames.size(); i++)
{
headerRow.createCell(i).setCellValue(columnNames.get(i));
}

// write content
for (int i = 0; i < _rows.size(); i++)
{
Map<String, Object> row = _rows.get(i);
SXSSFRow currentRow = sheet.createRow(i + 1);
for (int j = 0; j < columnNames.size(); j++)
{
currentRow.createCell(j).setCellValue(row.getOrDefault(columnNames.get(j), "").toString());
}
}
workbook.write(out);
}

return file;
}

/**
* Creates a file containing the contents of the current rows, formatted in TSV, CSV, or xlsx.
* The file is written to the test temp dir
Expand All @@ -712,33 +676,43 @@ public File writeGeneratedDataToExcel(String sheetName, String fileName) throws
*/
public File writeData(String fileName)
{
String fileExtension = fileName.toLowerCase().substring(fileName.lastIndexOf('.') + 1);
switch (fileExtension)
{
case "xlsx":
case "xls":
try
{
return writeGeneratedDataToExcel("sheet1", fileName);
}
catch (IOException e)
{
throw new RuntimeException(e);
}
case "csv":
return writeData(fileName, CSVFormat.DEFAULT);
case "tsv":
return writeData(fileName, CSVFormat.TDF);
default:
throw new IllegalArgumentException("Unsupported file extension: " + fileExtension);
}
return writeData(fileName, new FileRowIterator(getFieldsForFile(), _rows));
}

public File writeData(String fileName, CSVFormat format)
/**
* Generate rows and write to file without saving in data generator
* @param fileName the name of the file, e.g. 'testDataFileForMyTest.tsv'
* @param numRows the number of rows to generate
* @return object pointing at created file
*/
public File writeData(String fileName, int numRows)
{
return writeData(fileName, new FileRowIterator(getFieldsForFile(), this::generateRow, numRows));
}

/**
* Creates a file containing the contents of the current rows, formatted in TSV, CSV, or xlsx.
* The file is written to the test temp dir
* @param fileName the name of the file, e.g. 'testDataFileForMyTest.tsv'
* @return File object pointing at created file
*/
public File writeData(String fileName, Iterator<List<Object>> rowIterator)
{
try
{
return TestDataUtils.writeRowsToFile(fileName, TestDataUtils.rowListsFromMaps(_rows, getFieldsForFile()), format);
String fileExtension = fileName.toLowerCase().substring(fileName.lastIndexOf('.') + 1);
switch (fileExtension)
{
case "xlsx":
case "xls":
return TestDataUtils.writeRowsToExcel(fileName, rowIterator);
case "csv":
return TestDataUtils.writeRowsToFile(fileName, rowIterator, CSVFormat.DEFAULT);
case "tsv":
return TestDataUtils.writeRowsToFile(fileName, rowIterator, CSVFormat.TDF);
default:
throw new IllegalArgumentException("Unsupported file extension: " + fileExtension);
}
}
catch (IOException e)
{
Expand Down Expand Up @@ -915,3 +889,71 @@ public static boolean doesDomainExists(final String containerPath, final String
}

}

class FileRowIterator implements Iterator<List<Object>>
{
private final List<String> headers;
private final Iterator<Map<String, Object>> rows;

private boolean firstRow = true;

public FileRowIterator(@NotNull List<String> headers, @NotNull Iterator<Map<String, Object>> rows)
{
this.headers = Objects.requireNonNull(headers);
this.rows = Objects.requireNonNull(rows);
}

public FileRowIterator(@NotNull List<String> headers, @NotNull Supplier<Map<String, Object>> rowSupplier, final int rowCount)
{
this(headers, new Iterator<>()
{
int count = 0;

@Override
public boolean hasNext()
{
return count < rowCount;
}

@Override
public Map<String, Object> next()
{
count++;
return rowSupplier.get();
}
});
}

public FileRowIterator(@NotNull List<String> headers, @NotNull List<Map<String, Object>> rows)
{
this(headers, rows.iterator());
}

@Override
public boolean hasNext()
{
return firstRow || rows.hasNext();
}

@Override
public List<Object> next()
{
if (!hasNext())
throw new NoSuchElementException();

if (firstRow)
{
firstRow = false;
return Collections.unmodifiableList(headers);
}
else
{
return rowMapToList(rows.next());
}
}

private List<Object> rowMapToList(Map<String, Object> row)
{
return headers.stream().map(h -> row.getOrDefault(h, "")).toList();
}
}
64 changes: 62 additions & 2 deletions src/org/labkey/test/util/data/TestDataUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,14 +7,18 @@
import org.apache.commons.io.FileUtils;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.poi.xssf.streaming.SXSSFRow;
import org.apache.poi.xssf.streaming.SXSSFWorkbook;
import org.jetbrains.annotations.NotNull;
import org.labkey.serverapi.reader.DataLoader;
import org.labkey.serverapi.reader.TabLoader;
import org.labkey.test.TestFileUtils;
import org.labkey.test.params.FieldDefinition;
import org.labkey.test.util.TestDataGenerator;
import org.labkey.test.util.TestLogger;

import java.io.File;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
Expand All @@ -24,6 +28,7 @@
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
Expand Down Expand Up @@ -314,15 +319,70 @@ public static <T> File writeRowsToCsv(String fileName, List<List<T>> rows) throw
}

public static @NotNull <T> File writeRowsToFile(String fileName, List<List<T>> rows, CSVFormat format) throws IOException
{
return writeRowsToFile(fileName, rows.iterator(), format);
}

public static @NotNull <T> File writeRowsToFile(String fileName, Iterator<List<T>> rowIterator, CSVFormat format) throws IOException
{
File file = new File(TestFileUtils.getTestTempDir(), fileName);
FileUtils.forceMkdirParent(file);

TestLogger.log("Writing data file " + file.getAbsolutePath());

try (CSVPrinter printer = new CSVPrinter(new FileWriter(file, StandardCharsets.UTF_8), format)) {
for (List<T> row : rows)
while (rowIterator.hasNext())
{
printer.printRecord(row);
printer.printRecord(rowIterator.next());
}
}

return file;
}

public static @NotNull <T> File writeRowsToExcel(String fileName, List<List<T>> rows) throws IOException
{
return writeRowsToExcel(fileName, rows.iterator());
}

public static @NotNull <T> File writeRowsToExcel(String fileName, Iterator<List<T>> rowIterator) throws IOException
{
File file = new File(TestFileUtils.getTestTempDir(), fileName);
FileUtils.forceMkdirParent(file);

TestLogger.log("Writing data file " + file.getAbsolutePath());

try (SXSSFWorkbook workbook = new SXSSFWorkbook(1000); // only holds 1000 rows in memory
FileOutputStream out = new FileOutputStream(file))
{
var sheet = workbook.createSheet("sheet1");

// write headers as row 0
List<T> columnNames = rowIterator.next();
var headerRow = sheet.createRow(0);
for (int col = 0; col < columnNames.size(); col++)
{
if (columnNames.get(col) instanceof String s)
{
headerRow.createCell(col).setCellValue(s);
}
else
{
throw new IllegalArgumentException("Expected column headers to be Strings but got " + columnNames.get(col).getClass().getSimpleName());
}
}

// write content
for (int rowNum = 1; rowIterator.hasNext(); rowNum++)
{
List<T> row = rowIterator.next();
SXSSFRow currentRow = sheet.createRow(rowNum + 1);
for (int col = 0; col < columnNames.size(); col++)
{
currentRow.createCell(col).setCellValue(row.get(col).toString());
}
}
workbook.write(out);
}

return file;
Expand Down