Feature/cdspt 16114 add parquet streaming functionality

mario/data_extractor.py

@@ -3,10 +3,12 @@
 
 import pandas as pd
 from pandas import DataFrame
+import pyarrow as pa
+import pyarrow.parquet as pq
 
 from mario.dataset_specification import DatasetSpecification
 from mario.metadata import Metadata
-from mario.options import CsvOptions, HyperOptions
+from mario.options import CsvOptions, HyperOptions, OutputOptions
 from mario.mapping import FieldMapping
 
 logger = logging.getLogger(__name__)
@@ -458,6 +460,41 @@ def stream_sql_to_hyper(self, file_path: str, **kwargs):
                     row_counter += len(df)  # Update the counter
             frame_to_hyper(df, database=file_path, table=table_name, table_mode='a')
 
+    def stream_sql_to_parquet(self, file_path: str, **kwargs):
+        """
+        Write From SQL to parquet using streaming. No data is held in memory
+        apart from chunks of rows as they are read.
+        Optionally, data can be validated as it is read.
+        """
+        options = OutputOptions(**kwargs)
+        self.__build_query__()
+        logger.info("Executing query")
+
+        writer = None
+        connection = self.get_connection()
+        row_counter = 0
+        for df in pd.read_sql(self._query[0], connection, chunksize=options.chunk_size):
+            df = self.mapping.df_to_logical(df)
+            if options.validate or options.minimise or options.include_row_numbers:
+                self._data = df
+                if options.validate:
+                    self.validate_data(allow_nulls=options.allow_nulls)
+                if options.minimise:
+                    self.__minimise_data__()
+                    df = self._data
+                self._data = None
+                if options.include_row_numbers:
+                    df['row_number'] = range(row_counter, row_counter + len(df))
+                    row_counter += len(df)  # Update the counter
+
+            table = pa.Table.from_pandas(df)
+            if writer is None:
+                writer = pq.ParquetWriter(file_path, table.schema)
+
+            writer.write_table(table)
+
+        writer.close()
+
     def stream_sql_query_to_csv(self, file_path, query, connection, row_counter=0, **kwargs) -> int:
         from mario.query_builder import get_formatted_query
         options = CsvOptions(**kwargs)

test/test_data_extractor.py

@@ -162,6 +162,30 @@ def test_stream_sql_to_hyper():
     shutil.rmtree(folder.name)
 
 
+def test_stream_sql_to_parquet():
+    # Skip this test if we don't have a connection string
+    if not os.environ.get('CONNECTION_STRING'):
+        pytest.skip("Skipping SQL test as no database configured")
+
+    dataset = dataset_from_json(os.path.join('test', 'dataset.json'))
+    metadata = metadata_from_json(os.path.join('test', 'metadata.json'))
+    configuration = Configuration(
+        connection_string=os.environ.get('CONNECTION_STRING'),
+        schema="dev",
+        view="superstore",
+        query_builder=ViewBasedQueryBuilder
+    )
+    extractor = StreamingDataExtractor(
+        dataset_specification=dataset,
+        metadata=metadata,
+        configuration=configuration
+    )
+    file = tempfile.NamedTemporaryFile(suffix='.parquet')
+    extractor.stream_sql_to_parquet(file_path=file.name, chunk_size=1000)
+    df = pd.read_parquet(file)
+    assert len(df) == 10194
+
+
 def test_stream_sql_to_csv_with_validation():
     # Skip this test if we don't have a connection string
     if not os.environ.get('CONNECTION_STRING'):