fwrite supports complex-valued columns, part of #3690

NEWS.md

@@ -35,6 +35,8 @@
     * Gains `yaml` argument matching that of `fread`, [#3534](https://github.com/Rdatatable/data.table/issues/3534). See the item in `fread` for a bit more detail; here, we'd like to reiterate that feedback is appreciated in the initial phase of rollout for this feature.
 
     * Gains `bom` argument to add a *byte order mark* (BOM) at the beginning of the file to signal that the file is encoded in UTF-8, [#3488](https://github.com/Rdatatable/data.table/issues/3488). Thanks to Stefan Fleck for requesting and Philippe Chataignon for implementing.
+
+    * Supports writing complex-valued columns, part of [#3690](https://github.com/Rdatatable/data.table/issues/3690).
 
 4. Assigning to one item of a list column no longer requires the RHS to be wrapped with `list` or `.()`, [#950](https://github.com/Rdatatable/data.table/issues/950).
 

inst/tests/tests.Rraw

@@ -9566,8 +9566,6 @@ test(1658.44, file.info(f3)$size, file.info(f1)$size)
 unlink(c(f1,f2,f3))
 DT = data.table(a=1:3, b=list(1:4, c(3.14, 100e10), c("foo", "bar", "baz")))
 test(1658.45, fwrite(DT), output=c("a,b","1,1|2|3|4","2,3.14|1e+12","3,foo|bar|baz"))
-DT[3,b:=c(3i,4i,5i)]
-test(1658.46, fwrite(DT), error="Row 3 of list column is type 'complex'")
 DT[3,b:=factor(letters[1:3])]
 test(1658.47, fwrite(DT), error="Row 3 of list column is type 'factor'")
 
@@ -9595,6 +9593,14 @@ test(1658.53, fwrite(DT, file=f<-tempfile(fileext=".gz"), bom=TRUE, col.names=FA
 if (test_R.utils) test(1658.54, fread(f), setnames(DT,c("V1","V2")))
 unlink(f)
 
+# complex column support for fwrite, part of #3690
+DT = data.table(a = 1:3, z = 0:2 - (2:0)*1i)
+test(1658.54, fwrite(DT), output = '0-2i.*2,1-1i.*3,2+0i')
+test(1658.55, fwrite(data.table(exp(1) - pi*1i)), output = '2.718[0-9]*-3.141[0-9]*i')
+## formerly 1658.46
+DT = data.table(a=1:3, b=list(1:4, c(3.14, 100e10), c(3i,4i,5i)))
+test(1658.56, fwrite(DT), output='0+3i|0+4i|0+5i')
+
 ## End fwrite tests
 
 # tests for #679, inrange(), FR #707

src/fwrite.c

@@ -7,6 +7,7 @@
 #include <math.h>      // isfinite, isnan
 #include <stdlib.h>    // abs
 #include <string.h>    // strlen, strerror
+#include <complex.h>
 
 #ifdef WIN32
 #include <sys/types.h>
@@ -168,7 +169,7 @@ void genLookups() {
 }
 */
 
-void writeFloat64(double *col, int64_t row, char **pch)
+void writeFloat64Scalar(double x, char **pch)
 {
   // hand-rolled / specialized for speed
   // *pch is safely the output destination with enough space (ensured via calculating maxLineLen up front)
@@ -178,7 +179,6 @@ void writeFloat64(double *col, int64_t row, char **pch)
   //  ii) no C libary calls such as sprintf() where the fmt string has to be interpretted over and over
   // iii) no need to return variables or flags.  Just writes.
   //  iv) shorter, easier to read and reason with in one self contained place.
-  double x = col[row];
   char *ch = *pch;
   if (!isfinite(x)) {
     if (isnan(x)) {
@@ -292,6 +292,26 @@ void writeFloat64(double *col, int64_t row, char **pch)
   *pch = ch;
 }
 
+void writeFloat64(double *col, int64_t row, char **pch)
+{
+  double x = col[row];
+  char *ch = *pch;
+  writeFloat64Scalar(x, &ch);
+  *pch = ch;
+}
+
+void writeComplex(double complex *col, int64_t row, char **pch)
+{
+  double complex x = col[row];
+  char *ch = *pch;
+  writeFloat64Scalar(creal(x), &ch);
+  // let writeFloat64 handle the - sign for negative imaginary part
+  if (cimag(x) >= 0.0) *ch++ = '+';
+  writeFloat64Scalar(cimag(x), &ch);
+  *ch++ = 'i';
+  *pch = ch;
+}
+
 // DATE/TIME
 
 static inline void write_time(int32_t x, char **pch)

src/fwrite.h

@@ -9,12 +9,14 @@
 
 typedef void (*writer_fun_t)(void *, int64_t, char **);
 
+// in the order of writer_fun_t in fwriteR.c
 void writeBool8();
 void writeBool32();
 void writeBool32AsString();
 void writeInt32();
 void writeInt64();
 void writeFloat64();
+void writeComplex();
 void writeITime();
 void writeDateInt32();
 void writeDateFloat64();
@@ -33,6 +35,7 @@ typedef enum {   // same order as fun[] above
   WF_Int32,
   WF_Int64,
   WF_Float64,
+  WF_Complex,
   WF_ITime,
   WF_DateInt32,
   WF_DateFloat64,
@@ -50,6 +53,7 @@ static const int writerMaxLen[] = {  // same order as fun[] and WFs above; max f
   11, //&writeInt32            "-2147483647"
   20, //&writeInt64            "-9223372036854775807"
   29, //&writeFloat64          "-3.141592653589793115998E-123" [max sf 22 consistent with options()$digits]
+  60, //&writeComplex          "-3.141592653589793115998E-123+2.7182818284590450907956i" [3x writeFloat64,+,i]
   32, //&writeITime
   16, //&writeDateInt32
   16, //&writeDateFloat64

src/fwriteR.c

@@ -55,6 +55,7 @@ writer_fun_t funs[] = {
   &writeInt32,
   &writeInt64,
   &writeFloat64,
+  &writeComplex,
   &writeITime,
   &writeDateInt32,
   &writeDateFloat64,
@@ -130,6 +131,8 @@ static int32_t whichWriter(SEXP column) {
     if (INHERITS(column, char_Date))     return WF_DateFloat64;
     if (INHERITS(column, char_POSIXct))  return WF_POSIXct;
     return WF_Float64;
+  case CPLXSXP:
+    return WF_Complex;
   case STRSXP:
     return WF_String;
   case VECSXP: