SNOW-83333: Support arrow result format in chunk downloader.

Author: sfc-gh-stakeda

chunk_downloader.py

@@ -9,9 +9,20 @@
 from multiprocessing.pool import ThreadPool
 from threading import (Condition, Lock)
 
+from .compat import ITERATOR
 from snowflake.connector.network import ResultIterWithTimings
+from snowflake.connector.gzip_decoder import decompress_raw_data
+from snowflake.connector.util_text import split_rows_from_stream
 from .errorcode import (ER_NO_ADDITIONAL_CHUNK, ER_CHUNK_DOWNLOAD_FAILED)
 from .errors import (Error, OperationalError)
+import json
+from io import StringIO
+from gzip import GzipFile
+
+try:
+    from pyarrow.ipc import open_stream
+except ImportError:
+    pass
 
 DEFAULT_REQUEST_TIMEOUT = 3600
 
@@ -42,9 +53,11 @@ class SnowflakeChunkDownloader(object):
     """
 
     def _pre_init(self, chunks, connection, cursor, qrmk, chunk_headers,
+                  query_result_format='JSON',
                   prefetch_threads=DEFAULT_CLIENT_PREFETCH_THREADS,
                   use_ijson=False):
         self._use_ijson = use_ijson
+        self._query_result_format = query_result_format
 
         self._downloader_error = None
 
@@ -87,9 +100,11 @@ def _pre_init(self, chunks, connection, cursor, qrmk, chunk_headers,
         self._next_chunk_to_consume = 0
 
     def __init__(self, chunks, connection, cursor, qrmk, chunk_headers,
+                 query_result_format='JSON',
                  prefetch_threads=DEFAULT_CLIENT_PREFETCH_THREADS,
                  use_ijson=False):
         self._pre_init(chunks, connection, cursor, qrmk, chunk_headers,
+                       query_result_format=query_result_format,
                        prefetch_threads=prefetch_threads,
                        use_ijson=use_ijson)
         logger.debug('Chunk Downloader in memory')
@@ -251,10 +266,95 @@ def _fetch_chunk(self, url, headers):
         """
         Fetch the chunk from S3.
         """
+        handler = JsonBinaryHandler(is_raw_binary_iterator=True,
+                                    use_ijson=self._use_ijson) \
+            if self._query_result_format == 'json' else \
+            ArrowBinaryHandler()
+
         return self._connection.rest.fetch(
             u'get', url, headers,
             timeout=DEFAULT_REQUEST_TIMEOUT,
             is_raw_binary=True,
-            is_raw_binary_iterator=True,
-            use_ijson=self._use_ijson,
+            binary_data_handler=handler,
             return_timing_metrics=True)
+
+
+class RawBinaryDataHandler:
+    """
+    Abstract class being passed to network.py to handle raw binary data
+    """
+    def to_iterator(self, raw_data_fd):
+        pass
+
+
+class JsonBinaryHandler(RawBinaryDataHandler):
+    """
+    Convert result chunk in json format into interator
+    """
+    def __init__(self, is_raw_binary_iterator, use_ijson):
+        self._is_raw_binary_iterator = is_raw_binary_iterator
+        self._use_ijson = use_ijson
+
+    def to_iterator(self, raw_data_fd):
+        raw_data = decompress_raw_data(
+            raw_data_fd, add_bracket=True
+        ).decode('utf-8', 'replace')
+        if not self._is_raw_binary_iterator:
+            ret = json.loads(raw_data)
+        elif not self._use_ijson:
+            ret = iter(json.loads(raw_data))
+        else:
+            ret = split_rows_from_stream(StringIO(raw_data))
+        return ret
+
+
+class ArrowBinaryHandler(RawBinaryDataHandler):
+    """
+    Handler to consume data as arrow stream
+    """
+    def to_iterator(self, raw_data_fd):
+        gzip_decoder = GzipFile(fileobj=raw_data_fd, mode='r')
+        reader = open_stream(gzip_decoder)
+        return ArrowChunkIterator(reader)
+
+
+class ArrowChunkIterator(ITERATOR):
+    """
+    Given a list of record batches, iterate over
+    these batches row by row
+    """
+
+    def __init__(self, arrow_stream_reader):
+        self._batches = []
+        for record_batch in arrow_stream_reader:
+            self._batches.append(record_batch.columns)
+
+        self._column_count = len(self._batches[0])
+        self._batch_count = len(self._batches)
+        self._batch_index = -1
+        self._index_in_batch = -1
+        self._row_count_in_batch = 0
+
+    def next(self):
+        return self.__next__()
+
+    def __next__(self):
+        self._index_in_batch += 1
+        if self._index_in_batch < self._row_count_in_batch:
+            return self._return_row()
+        else:
+            self._batch_index += 1
+            if self._batch_index < self._batch_count:
+                self._index_in_batch = 0
+                self._row_count_in_batch = len(self._batches[self._batch_index][0])
+                return self._return_row()
+
+        raise StopIteration
+
+    def _return_row(self):
+        ret = []
+        current_batch = self._batches[self._batch_index]
+        for col_array in current_batch:
+            ret.append(col_array[self._index_in_batch])
+
+        return ret

converter_arrow.py

@@ -0,0 +1,73 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright (c) 2012-2019 Snowflake Computing Inc. All right reserved.
+#
+
+from logging import getLogger
+from decimal import Context
+from datetime import datetime, timedelta, date
+from .converter import SnowflakeConverter
+
+logger = getLogger(__name__)
+
+ZERO_EPOCH = datetime.utcfromtimestamp(0)
+
+
+class SnowflakeArrowConverter(SnowflakeConverter):
+    """
+    Convert from arrow data into python native data types
+    """
+
+    def to_python_method(self, type_name, column):
+        ctx = column.copy()
+
+        if type_name == 'FIXED' and ctx['scale'] != 0:
+            ctx['decimalCtx'] = Context(prec=ctx['precision'])
+
+        converters = [u'_{type_name}_to_python'.format(type_name=type_name)]
+        if self._use_numpy:
+            converters.insert(0, u'_{type_name}_numpy_to_python'.format(
+                type_name=type_name))
+        for conv in converters:
+            try:
+                return getattr(self, conv)(ctx)
+            except AttributeError:
+                pass
+        logger.warning(
+            "No column converter found for type: %s", type_name)
+        return None  # Skip conversion
+
+    def _FIXED_to_python(self, ctx):
+        if ctx['scale'] == 0:
+            return lambda x: x.as_py()
+        else:
+            return lambda x, decimal_ctx=ctx['decimalCtx']: decimal_ctx.create_decimal(x.as_py())
+
+    def _REAL_to_python(self, _):
+        return lambda x: x.as_py()
+
+    def _TEXT_to_python(self, _):
+        return lambda x: x.as_py()
+
+    def _BINARY_to_python(self, _):
+        return lambda x: x.as_py()
+
+    def _VARIANT_to_python(self, _):
+        return lambda x: x.as_py()
+
+    def _BOOLEAN_to_python(self, _):
+        return lambda x: x.as_py() > 0
+
+    def _DATE_to_python(self, _):
+
+        def conv(value):
+            try:
+                return datetime.utcfromtimestamp(value * 86400).date()
+            except OSError as e:
+                logger.debug("Failed to convert: %s", e)
+                ts = ZERO_EPOCH + timedelta(
+                    seconds=value * (24 * 60 * 60))
+                return date(ts.year, ts.month, ts.day)
+
+        return conv

cursor.py

@@ -10,7 +10,7 @@
 import uuid
 from logging import getLogger
 from threading import (Timer, Lock)
-
+from base64 import b64decode
 from six import u
 
 from .compat import (BASE_EXCEPTION_CLASS)
@@ -30,6 +30,13 @@
 from .sqlstate import (SQLSTATE_FEATURE_NOT_SUPPORTED)
 from .telemetry import (TelemetryData, TelemetryField)
 from .time_util import get_time_millis
+from .chunk_downloader import ArrowChunkIterator
+from .converter_arrow import SnowflakeArrowConverter
+
+try:
+    from pyarrow.ipc import open_stream
+except ImportError:
+    pass
 
 STATEMENT_TYPE_ID_DML = 0x3000
 STATEMENT_TYPE_ID_INSERT = STATEMENT_TYPE_ID_DML + 0x100
@@ -570,6 +577,7 @@ def _is_dml(self, data):
 
     def chunk_info(self, data, use_ijson=False):
         is_dml = self._is_dml(data)
+        self._query_result_format = data.get(u'queryResultFormat', u'json')
 
         if self._total_rowcount == -1 and not is_dml and data.get(u'total') \
                 is not None:
@@ -578,6 +586,10 @@ def chunk_info(self, data, use_ijson=False):
         self._description = []
         self._column_idx_to_name = {}
         self._column_converter = []
+
+        converter = SnowflakeArrowConverter() if \
+            self._query_result_format == 'arrow' else self._connection.converter
+
         for idx, column in enumerate(data[u'rowtype']):
             self._column_idx_to_name[idx] = column[u'name']
             type_value = FIELD_NAME_TO_ID[column[u'type'].upper()]
@@ -589,15 +601,21 @@ def chunk_info(self, data, use_ijson=False):
                                       column[u'scale'],
                                       column[u'nullable']))
             self._column_converter.append(
-                self._connection.converter.to_python_method(
-                    column[u'type'].upper(), column))
+                    converter.to_python_method(
+                        column[u'type'].upper(), column))
 
         self._total_row_index = -1  # last fetched number of rows
 
         self._chunk_index = 0
         self._chunk_count = 0
-        self._current_chunk_row = iter(data.get(u'rowset'))
-        self._current_chunk_row_count = len(data.get(u'rowset'))
+        if self._query_result_format == 'arrow':
+            # result as arrow chunk
+            arrow_bytes = b64decode(data.get(u'rowsetBase64'))
+            arrow_reader = open_stream(arrow_bytes)
+            self._current_chunk_row = ArrowChunkIterator(arrow_reader)
+        else:
+            self._current_chunk_row = iter(data.get(u'rowset'))
+            self._current_chunk_row_count = len(data.get(u'rowset'))
 
         if u'chunks' in data:
             chunks = data[u'chunks']
@@ -619,6 +637,7 @@ def chunk_info(self, data, use_ijson=False):
             logger.debug(u'qrmk=%s', qrmk)
             self._chunk_downloader = self._connection._chunk_downloader_class(
                 chunks, self._connection, self, qrmk, chunk_headers,
+                query_result_format=self._query_result_format,
                 prefetch_threads=self._connection.client_prefetch_threads,
                 use_ijson=use_ijson)
 

network.py

@@ -13,7 +13,7 @@
 import sys
 import time
 import uuid
-from io import StringIO, BytesIO
+from io import BytesIO
 from threading import Lock
 
 import OpenSSL.SSL
@@ -47,7 +47,6 @@
                      InterfaceError, InternalServerError, ForbiddenError,
                      BadGatewayError, BadRequest, MethodNotAllowed,
                      OtherHTTPRetryableError)
-from .gzip_decoder import decompress_raw_data
 from .sqlstate import (SQLSTATE_CONNECTION_NOT_EXISTS,
                        SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED,
                        SQLSTATE_CONNECTION_REJECTED)
@@ -56,7 +55,6 @@
     DEFAULT_MASTER_VALIDITY_IN_SECONDS
 )
 from .tool.probe_connection import probe_connection
-from .util_text import split_rows_from_stream
 from .version import VERSION
 
 if PY2:
@@ -739,8 +737,7 @@ def _request_exec(
             catch_okta_unauthorized_error=False,
             is_raw_text=False,
             is_raw_binary=False,
-            is_raw_binary_iterator=True,
-            use_ijson=False,
+            binary_data_handler=None,
             socket_timeout=DEFAULT_SOCKET_CONNECT_TIMEOUT,
             return_timing_metrics=False):
         if socket_timeout > DEFAULT_SOCKET_CONNECT_TIMEOUT:
@@ -785,15 +782,7 @@ def _request_exec(
                         ret = raw_ret.text
                     elif is_raw_binary:
                         start_time = get_time_millis()
-                        raw_data = decompress_raw_data(
-                            raw_ret.raw, add_bracket=True
-                        ).decode('utf-8', 'replace')
-                        if not is_raw_binary_iterator:
-                            ret = json.loads(raw_data)
-                        elif not use_ijson:
-                            ret = iter(json.loads(raw_data))
-                        else:
-                            ret = split_rows_from_stream(StringIO(raw_data))
+                        ret = binary_data_handler.to_iterator(raw_ret.raw)
                         timing_metrics[
                             ResultIterWithTimings.PARSE] = get_time_millis() - start_time
 

scripts/install.bat

@@ -7,6 +7,7 @@ call env\Scripts\activate
 # https://github.com/pypa/pip/issues/6566
 python -m pip install --upgrade pip==18.1
 pip install pendulum
+pip install pyarrow
 pip install numpy
 pip install pytest pytest-cov pytest-rerunfailures
 pip install .

scripts/install.sh

@@ -31,6 +31,7 @@ fi
 
 source ./venv/bin/activate
 pip install numpy pendulum
+pip install pyarrow
 pip install pytest pytest-cov pytest-rerunfailures
 if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]] || [[ $PYTHON_VERSION == "2.7"* ]]; then
     pip install mock

setup.py

@@ -85,6 +85,10 @@
         "secure-local-storage": [
             'keyring!=16.1.0'
         ],
+        "arrow-result": [
+            'pyarrow>=0.13.0;python_version>"3.4"',
+            'pyarrow>=0.13.0;python_version<"3.0"'
+        ]
     },
 
     classifiers=[

test/conftest.py

@@ -126,6 +126,8 @@ def get_db_parameters():
     ret['name_wh'] = ret['name'] + 'wh'
 
     ret['schema'] = TEST_SCHEMA
+
+    # This reduces a chance to exposing password in test output.
     ret['a00'] = 'dummy parameter'
     ret['a01'] = 'dummy parameter'
     ret['a02'] = 'dummy parameter'