kuzudb · aziz-mu · Nov 18, 2022 · Nov 17, 2022
diff --git a/src/main/include/query_result.h b/src/main/include/query_result.h
@@ -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(); }
 

diff --git a/src/main/query_result.cpp b/src/main/query_result.cpp
@@ -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++) {
@@ -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;
             }
         }
     }

diff --git a/test/main/csv_output_test.cpp b/test/main/csv_output_test.cpp
@@ -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());
+}
diff --git a/tools/python_api/include/py_query_result.h b/tools/python_api/include/py_query_result.h
@@ -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();
 

diff --git a/tools/python_api/py_query_result.cpp b/tools/python_api/py_query_result.cpp
@@ -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"
@@ -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)
@@ -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() {

diff --git a/tools/python_api/test/test_write_to_csv.py b/tools/python_api/test/test_write_to_csv.py
@@ -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]"
 """ 
@@ -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'']"])