Feature/csv destination

earthmover/nodes/destination.py

@@ -1,4 +1,5 @@
 import csv
+import dask.dataframe as dd
 import jinja2
 import os
 import pandas as pd
@@ -18,8 +19,11 @@ class Destination(Node):
     mode: str = None  # Documents which class was chosen.
     allowed_configs: Tuple[str] = ('debug', 'expect', 'show_progress', 'repartition', 'source',)
 
-    def __new__(cls, *args, **kwargs):
-        return object.__new__(FileDestination)
+    def __new__(cls, name: str, config: 'YamlMapping', *, earthmover: 'Earthmover'):
+        if config.get('extension') == 'csv' or config.get('extension') == 'tsv':
+            return object.__new__(CsvDestination)
+        else:
+            return object.__new__(FileDestination)
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -132,3 +136,79 @@ def render_row(self, row: pd.Series):
             raise
 
         return json_string
+
+
+class CsvDestination(Destination):
+    """
+
+    """
+    mode: str = 'csv'  # Documents which class was chosen.
+    allowed_configs: Tuple[str] = (
+        'debug', 'expect', 'show_progress', 'repartition', 'source',
+        'extension', 'header', 'separator', 'limit', 'keep_columns'
+    )
+
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.header = self.error_handler.assert_get_key(self.config, 'header', dtype=bool, required=False, default=True)
+        self.separator = self.error_handler.assert_get_key(self.config, 'separator', dtype=str, required=False, default=",")
+        self.limit = self.error_handler.assert_get_key(self.config, 'limit', dtype=int, required=False, default=None)
+        self.extension = self.error_handler.assert_get_key(self.config, 'extension', dtype=str, required=False, default="csv")
+        self.keep_columns = self.error_handler.assert_get_key(self.config, 'keep_columns', required=False, default=None)
+
+        self.file = os.path.join(
+            self.earthmover.state_configs['output_dir'],
+            f"{self.name}.{self.extension}"
+        )   
+
+    def execute(self, **kwargs):
+        """
+
+        :return:
+        """
+        super().execute(**kwargs)
+
+        self.data = self.upstream_sources[self.source].data
+
+        # Apply limit to dataframe if specified.
+        if self.limit:
+            if self.limit > len(self.data):
+                self.error_handler.throw(
+                    f"Limit value exceeds the number of rows in the data"
+                )
+                raise 
+
+            self.data = dd.from_pandas(self.data.head(n=self.limit), npartitions=1)
+
+        # Verify the output directory exists.
+        os.makedirs(os.path.dirname(self.file), exist_ok=True)
+        self.logger.info(f"Directory created: {os.path.dirname(self.file)}")
+
+        # Subset dataframe columns if specified
+        try:
+            if self.keep_columns:
+                self.data = self.data[self.keep_columns]
+
+        except KeyError as e:
+            self.error_handler.throw(
+                f"Error occurred while subsetting the data: {e.args[0]}"
+            )
+            raise
+
+        # Change separator to tab if extension is tsv
+        if self.extension == 'tsv':
+            self.separator = '\t'
+
+        try:
+            self.data.to_csv(
+                filename=self.file, single_file=True, index=False,
+                sep=self.separator, header=self.header
+            )
+            self.logger.info(f"Output `{self.file}` written")
+
+        except Exception as err:
+            self.error_handler.throw(
+                f"Error writing data to {self.extension} file: ({err})"
+            )
+            raise

example_projects/01_simple/earthmover.yaml

@@ -50,4 +50,14 @@ destinations:
     source: $sources.testing_source
     template: ./templates/studentSchoolAttendanceEvent.jsont
     extension: jsonl
-    linearize: True
+    linearize: True
+  attendance:
+    source: $sources.attendance
+    extension: tsv
+    debug: True
+    header: True
+    limit: 5
+    keep_columns: 
+      - day
+      - student_id
+