Skip to content

Commit

Permalink
added custom specifiers to writeToCSV
Browse files Browse the repository at this point in the history
  • Loading branch information
aziz-mu committed Nov 18, 2022
1 parent 3aff141 commit 9282aa0
Show file tree
Hide file tree
Showing 6 changed files with 90 additions and 18 deletions.
3 changes: 2 additions & 1 deletion src/main/include/query_result.h
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,8 @@ class QueryResult {
// TODO: this is not efficient and should be replaced by iterator
std::shared_ptr<processor::FlatTuple> getNext();

void writeToCSV(string fileName);
void writeToCSV(
string fileName, char delimiter = ',', char escapeCharacter = '"', char newline = '\n');

inline uint64_t getNumColumns() const { return header->columnDataTypes.size(); }

Expand Down
38 changes: 25 additions & 13 deletions src/main/query_result.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -31,10 +31,12 @@ shared_ptr<FlatTuple> QueryResult::getNext() {
return iterator->getNextFlatTuple();
}

void QueryResult::writeToCSV(string fileName) {
void QueryResult::writeToCSV(string fileName, char delimiter, char escapeCharacter, char newline) {
ofstream file;
file.open(fileName);
shared_ptr<FlatTuple> nextTuple;
assert(delimiter != '\0');
assert(newline != '\0');
while (hasNext()) {
nextTuple = getNext();
for (auto idx = 0ul; idx < nextTuple->len(); idx++) {
Expand All @@ -45,33 +47,43 @@ void QueryResult::writeToCSV(string fileName) {
isStringList = true;
}
bool surroundQuotes = false;
string thisValueStr;
string csvStr;
for (long unsigned int j = 0; j < resultVal.length(); j++) {
if (!surroundQuotes) {
if (resultVal[j] == '"' || resultVal[j] == '\n' || resultVal[j] == ',') {
if (resultVal[j] == escapeCharacter || resultVal[j] == newline ||
resultVal[j] == delimiter) {
surroundQuotes = true;
}
}
if (resultVal[j] == '"') {
thisValueStr += "\"\"";
if (resultVal[j] == escapeCharacter) {
csvStr += escapeCharacter;
csvStr += escapeCharacter;
} else if (resultVal[j] == ',' && isStringList) {
thisValueStr += "\"\",\"\"";
csvStr += escapeCharacter;
csvStr += escapeCharacter;
csvStr += ',';
csvStr += escapeCharacter;
csvStr += escapeCharacter;
} else if (resultVal[j] == '[' && isStringList) {
thisValueStr += "[\"\"";
csvStr += "[";
csvStr += escapeCharacter;
csvStr += escapeCharacter;
} else if (resultVal[j] == ']' && isStringList) {
thisValueStr += "\"\"]";
csvStr += escapeCharacter;
csvStr += escapeCharacter;
csvStr += "]";
} else {
thisValueStr += resultVal[j];
csvStr += resultVal[j];
}
}
if (surroundQuotes) {
thisValueStr = '"' + thisValueStr + '"';
csvStr = escapeCharacter + csvStr + escapeCharacter;
}
file << thisValueStr;
file << csvStr;
if (idx < nextTuple->len() - 1) {
file << ",";
file << delimiter;
} else {
file << endl;
file << newline;
}
}
}
Expand Down
29 changes: 29 additions & 0 deletions test/main/csv_output_test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -43,3 +43,32 @@ TEST_F(CSVOutputTest, ListCSVTest) {
string fileString = ss.str();
ASSERT_STREQ(fileString.c_str(), listOutput.c_str());
}

TEST_F(CSVOutputTest, AlternateDelimCSVTest) {
string listOutput = R"(ABFsUni "-2"-CsWork "-100"-DEsWork 7-)";
auto query = "MATCH (o:organisation) RETURN o.name, o.score";
auto result = conn->query(query);
result->writeToCSV("output_CSV_LIST.csv", '\t', '"', '-');
ifstream f("output_CSV_LIST.csv");
ostringstream ss;
ss << f.rdbuf();
string fileString = ss.str();
ASSERT_STREQ(fileString.c_str(), listOutput.c_str());
}

TEST_F(CSVOutputTest, AlternateEscapeCSVTest) {
string newline = "\n";
string listOutput =
R"(`[10,5]`,Alice)" + newline + R"(`[12,8]`,Bob)" + newline + R"(`[4,5]`,Carol)" + newline +
R"(`[1,9]`,Dan)" + newline + R"([2],Elizabeth)" + newline + R"(`[3,4,5,6,7]`,Farooq)" +
newline + R"([1],Greg)" + newline +
R"(`[10,11,12,3,4,5,6,7]`,Hubert Blaine Wolfeschlegelsteinhausenbergerdorff)" + newline;
auto query = "MATCH (p:person) RETURN p.workedHours, p.fName";
auto result = conn->query(query);
result->writeToCSV("output_CSV_LIST.csv", ',', '`');
ifstream f("output_CSV_LIST.csv");
ostringstream ss;
ss << f.rdbuf();
string fileString = ss.str();
ASSERT_STREQ(fileString.c_str(), listOutput.c_str());
}
2 changes: 1 addition & 1 deletion tools/python_api/include/py_query_result.h
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ class PyQueryResult {

py::list getNext();

void writeToCSV(py::str filename);
void writeToCSV(py::str filename, py::str delimiter, py::str escapeCharacter, py::str newline);

void close();

Expand Down
16 changes: 13 additions & 3 deletions tools/python_api/py_query_result.cpp
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
#include "include/py_query_result.h"

#include <fstream>
#include <string>

#include "datetime.h" // python lib
#include "include/py_query_result_converter.h"
Expand All @@ -11,7 +12,9 @@ void PyQueryResult::initialize(py::handle& m) {
py::class_<PyQueryResult>(m, "result")
.def("hasNext", &PyQueryResult::hasNext)
.def("getNext", &PyQueryResult::getNext)
.def("writeToCSV", &PyQueryResult::writeToCSV)
.def("writeToCSV", &PyQueryResult::writeToCSV, py::arg("filename"),
py::arg("delimiter") = ",", py::arg("escapeCharacter") = "\"",
py::arg("newline") = "\n")
.def("close", &PyQueryResult::close)
.def("getAsDF", &PyQueryResult::getAsDF)
.def("getColumnNames", &PyQueryResult::getColumnNames)
Expand All @@ -35,8 +38,15 @@ py::list PyQueryResult::getNext() {
return move(result);
}

void PyQueryResult::writeToCSV(py::str filename) {
queryResult->writeToCSV(filename);
void PyQueryResult::writeToCSV(
py::str filename, py::str delimiter, py::str escapeCharacter, py::str newline) {
std::string delimiterStr = delimiter;
std::string escapeCharacterStr = escapeCharacter;
std::string newlineStr = newline;
assert(delimiterStr.size() == 1);
assert(escapeCharacterStr.size() == 1);
assert(newlineStr.size() == 1);
queryResult->writeToCSV(filename, delimiterStr[0], escapeCharacterStr[0], newlineStr[0]);
}

void PyQueryResult::close() {
Expand Down
20 changes: 20 additions & 0 deletions tools/python_api/test/test_write_to_csv.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from pandas import read_csv

def test_write_to_csv(establish_connection):
outputString = """Alice,[""Aida""],"[10,5]"\nBob,[""Bobby""],"[12,8]"\nCarol,"[""Carmen"",""Fred""]","[4,5]"\nDan,"[""Wolfeschlegelstein"",""Daniel""]","[1,9]"\nElizabeth,[""Ein""],[2]\nFarooq,[""Fesdwe""],"[3,4,5,6,7]"\nGreg,[""Grad""],[1]\nHubert Blaine Wolfeschlegelsteinhausenbergerdorff,"[""Ad"",""De"",""Hi"",""Kye"",""Orlan""]","[10,11,12,3,4,5,6,7]"
"""
Expand All @@ -7,3 +9,21 @@ def test_write_to_csv(establish_connection):
with open("test_PYTHON_CSV.csv") as csv_file:
data = csv_file.read()
assert(data == outputString)

def test_write_to_csv_extra_args(establish_connection):
outputString = """35|1~30|2~45|1~20|2~20|1~25|2~40|2~83|2~"""
conn, db = establish_connection
result = conn.execute("MATCH (a:person) RETURN a.age, a.gender")
result.writeToCSV("test_PYTHON_CSV.csv", "|", '"', '~')
with open("test_PYTHON_CSV.csv") as csv_file:
data = csv_file.read()
assert(data == outputString)

def test_pandas_read_csv_extra_args(establish_connection):
conn, db = establish_connection
result = conn.execute("MATCH (a:person) RETURN a.fName, a.workedHours, a.usedNames");
result.writeToCSV("test_PYTHON_CSV.csv", "|", "'", "~");
df = read_csv("test_PYTHON_CSV.csv", delimiter = "|", lineterminator = "~", escapechar = "`")
assert(df.iloc[:,0].tolist() == ['Bob', 'Carol', 'Dan', 'Elizabeth', 'Farooq', 'Greg', 'Hubert Blaine Wolfeschlegelsteinhausenbergerdorff'])
assert(df.iloc[:,1].tolist() == ['[12,8]', '[4,5]', '[1,9]', '[2]', '[3,4,5,6,7]','[1]','[10,11,12,3,4,5,6,7]'])
assert(df.iloc[:,2].tolist() == ["[''Bobby'']", "[''Carmen'',''Fred'']", "[''Wolfeschlegelstein'',''Daniel'']","[''Ein'']","[''Fesdwe'']","[''Grad'']","[''Ad'',''De'',''Hi'',''Kye'',''Orlan'']"])

0 comments on commit 9282aa0

Please sign in to comment.