vastdb 0.1.1__py3-none-any.whl → 0.1.2__py3-none-any.whl

vastdb/internal_commands.py

@@ -8,7 +8,7 @@ import urllib.parse
 from collections import defaultdict, namedtuple
 from enum import Enum
 from ipaddress import IPv4Address, IPv6Address
-from typing import Iterator, Optional, Union
+from typing import Any, Dict, Iterator, List, Optional, Union
 
 import flatbuffers
 import ibis
@@ -92,7 +92,7 @@ UINT64_MAX = 18446744073709551615
 TABULAR_KEEP_ALIVE_STREAM_ID = 0xFFFFFFFF
 TABULAR_QUERY_DATA_COMPLETED_STREAM_ID = 0xFFFFFFFF - 1
 TABULAR_QUERY_DATA_FAILED_STREAM_ID = 0xFFFFFFFF - 2
-TABULAR_INVALID_ROW_ID = 0xFFFFFFFFFFFF # (1<<48)-1
+TABULAR_INVALID_ROW_ID = 0xFFFFFFFFFFFF  # (1<<48)-1
 ESTORE_INVALID_EHANDLE = UINT64_MAX
 IMPORTED_OBJECTS_TABLE_NAME = "vastdb-imported-objects"
 
@@ -127,11 +127,11 @@ def get_unit_to_flatbuff_time_unit(type):
     }
     return unit_to_flatbuff_time_unit[type]
 
+
 class Predicate:
     def __init__(self, schema: 'pa.Schema', expr: ibis.expr.types.BooleanColumn):
         self.schema = schema
         self.expr = expr
-        self.builder = None
 
     def get_field_indexes(self, field: 'pa.Field', field_name_per_index: list) -> None:
         field_name_per_index.append(field.name)
@@ -157,8 +157,8 @@ class Predicate:
             self._field_name_per_index = {field: index for index, field in enumerate(_field_name_per_index)}
         return self._field_name_per_index
 
-    def get_projections(self, builder: 'flatbuffers.builder.Builder', field_names: list = None):
-        if not field_names:
+    def get_projections(self, builder: 'flatbuffers.builder.Builder', field_names: Optional[List[str]] = None):
+        if field_names is None:
             field_names = self.field_name_per_index.keys()
         projection_fields = []
         for field_name in field_names:
@@ -172,7 +172,11 @@ class Predicate:
         return builder.EndVector()
 
     def serialize(self, builder: 'flatbuffers.builder.Builder'):
-        from ibis.expr.operations.generic import IsNull, Literal, TableColumn
+        from ibis.expr.operations.generic import (
+            IsNull,
+            Literal,
+            TableColumn,
+        )
         from ibis.expr.operations.logical import (
             And,
             Equals,
@@ -198,7 +202,7 @@ class Predicate:
             StringContains: self.build_match_substring,
         }
 
-        positions_map = dict((f.name, index) for index, f in enumerate(self.schema)) # TODO: BFS
+        positions_map = dict((f.name, index) for index, f in enumerate(self.schema))  # TODO: BFS
 
         self.builder = builder
 
@@ -215,7 +219,7 @@ class Predicate:
                 prev_field_name = None
                 for inner_op in or_args:
                     _logger.debug('inner_op %s', inner_op)
-                    builder_func = builder_map.get(type(inner_op))
+                    builder_func: Any = builder_map.get(type(inner_op))
                     if not builder_func:
                         raise NotImplementedError(inner_op.name)
 
@@ -270,20 +274,6 @@ class Predicate:
         fb_expression.AddImpl(self.builder, ref)
         return fb_expression.End(self.builder)
 
-    def build_domain(self, column: int, field_name: str):
-        offsets = []
-        filters = self.filters[field_name]
-        if not filters:
-            return self.build_or([self.build_is_not_null(column)])
-
-        field_name, *field_attrs = field_name.split('.')
-        field = self.schema.field(field_name)
-        for attr in field_attrs:
-            field = field.type[attr]
-        for filter_by_name in filters:
-            offsets.append(self.build_range(column=column, field=field, filter_by_name=filter_by_name))
-        return self.build_or(offsets)
-
     def rule_to_operator(self, raw_rule: str):
         operator_matcher = {
             'eq': self.build_equal,
@@ -339,6 +329,8 @@ class Predicate:
         return fb_expression.End(self.builder)
 
     def build_literal(self, field: pa.Field, value):
+        literal_type: Any
+
         if field.type.equals(pa.int64()):
             literal_type = fb_int64_lit
             literal_impl = LiteralImpl.Int64Literal
@@ -551,13 +543,20 @@ class Predicate:
         return self.build_function('match_substring', column, literal)
 
 
+class FieldNodesState:
+    def __init__(self) -> None:
+        # will be set during by the parser (see below)
+        self.buffers: Dict[int, Any] = defaultdict(lambda: None)  # a list of Arrow buffers (https://arrow.apache.org/docs/format/Columnar.html#buffer-listing-for-each-layout)
+        self.length: Dict[int, Any] = defaultdict(lambda: None)  # each array must have it's length specified (https://arrow.apache.org/docs/python/generated/pyarrow.Array.html#pyarrow.Array.from_buffers)
+
+
 class FieldNode:
     """Helper class for representing nested Arrow fields and handling QueryData requests"""
     def __init__(self, field: pa.Field, index_iter, parent: Optional['FieldNode'] = None, debug: bool = False):
-        self.index = next(index_iter) # we use DFS-first enumeration for communicating the column positions to VAST
+        self.index = next(index_iter)  # we use DFS-first enumeration for communicating the column positions to VAST
         self.field = field
         self.type = field.type
-        self.parent = parent # will be None if this is the top-level field
+        self.parent = parent  # will be None if this is the top-level field
         self.debug = debug
         if isinstance(self.type, pa.StructType):
             self.children = [FieldNode(field, index_iter, parent=self) for field in self.type]
@@ -574,11 +573,7 @@ class FieldNode:
             field = pa.field('entries', pa.struct([self.type.key_field, self.type.item_field]))
             self.children = [FieldNode(field, index_iter, parent=self)]
         else:
-            self.children = [] # for non-nested types
-
-        # will be set during by the parser (see below)
-        self.buffers = None # a list of Arrow buffers (https://arrow.apache.org/docs/format/Columnar.html#buffer-listing-for-each-layout)
-        self.length = None # each array must have it's length specified (https://arrow.apache.org/docs/python/generated/pyarrow.Array.html#pyarrow.Array.from_buffers)
+            self.children = []  # for non-nested types
 
     def _iter_to_root(self) -> Iterator['FieldNode']:
         yield self
@@ -599,22 +594,14 @@ class FieldNode:
             for child in self.children:
                 yield from child._iter_leaves()
 
-    def _iter_leaves(self) -> Iterator['FieldNode']:
-        """Generate only leaf nodes (i.e. columns having scalar types)."""
-        if not self.children:
-            yield self
-        else:
-            for child in self.children:
-                yield from child._iter_leaves()
-
     def debug_log(self, level=0):
         """Recursively dump this node state to log."""
         bufs = self.buffers and [b and b.hex() for b in self.buffers]
-        _logger.debug('%s%d: %s, bufs=%s, len=%s', '    '*level, self.index, self.field, bufs, self.length)
+        _logger.debug('%s%d: %s, bufs=%s, len=%s', '    ' * level, self.index, self.field, bufs, self.length)
         for child in self.children:
-            child.debug_log(level=level+1)
+            child.debug_log(level=level + 1)
 
-    def set(self, arr: pa.Array):
+    def set(self, arr: pa.Array, state: FieldNodesState):
         """
         Assign the relevant Arrow buffers from the received array into this node.
 
@@ -626,34 +613,39 @@ class FieldNode:
         For example, `Struct<A, B>` is sent as two separate columns: `Struct<A>` and `Struct<B>`.
         Also, `Map<K, V>` is sent (as its underlying representation): `List<Struct<K>>` and `List<Struct<V>>`
         """
-        buffers = arr.buffers()[:arr.type.num_buffers] # slicing is needed because Array.buffers() returns also nested array buffers
+        buffers = arr.buffers()[:arr.type.num_buffers]  # slicing is needed because Array.buffers() returns also nested array buffers
         if self.debug:
             _logger.debug("set: index=%d %s %s", self.index, self.field, [b and b.hex() for b in buffers])
-        if self.buffers is None:
-            self.buffers = buffers
-            self.length = len(arr)
+        if state.buffers[self.index] is None:
+            state.buffers[self.index] = buffers
+            state.length[self.index] = len(arr)
         else:
             # Make sure subsequent assignments are consistent with each other
             if self.debug:
-                if not self.buffers == buffers:
-                    raise ValueError(f'self.buffers: {self.buffers} are not equal with buffers: {buffers}')
-            if not self.length == len(arr):
-                raise ValueError(f'self.length: {self.length} are not equal with len(arr): {len(arr)}')
+                if not state.buffers[self.index] == buffers:
+                    raise ValueError(f'self.buffers: {state.buffers[self.index]} are not equal with buffers: {buffers}')
+            if not state.length[self.index] == len(arr):
+                raise ValueError(f'self.length: {state.length[self.index]} are not equal with len(arr): {len(arr)}')
 
-    def build(self) -> pa.Array:
+    def build(self, state: FieldNodesState) -> pa.Array:
         """Construct an Arrow array from the collected buffers (recursively)."""
-        children = self.children and [node.build() for node in self.children]
-        result = pa.Array.from_buffers(self.type, self.length, buffers=self.buffers, children=children)
+        children = self.children and [node.build(state) for node in self.children]
+        result = pa.Array.from_buffers(self.type, state.length[self.index], buffers=state.buffers[self.index], children=children)
         if self.debug:
             _logger.debug('%s result=%s', self.field, result)
         return result
 
 
 class QueryDataParser:
+    class QueryDataParserState(FieldNodesState):
+        def __init__(self) -> None:
+            super().__init__()
+            self.leaf_offset = 0
+
     """Used to parse VAST QueryData RPC response."""
     def __init__(self, arrow_schema: pa.Schema, *, debug=False):
         self.arrow_schema = arrow_schema
-        index = itertools.count() # used to generate leaf column positions for VAST QueryData RPC
+        index = itertools.count()  # used to generate leaf column positions for VAST QueryData RPC
         self.nodes = [FieldNode(field, index, debug=debug) for field in arrow_schema]
         self.debug = debug
         if self.debug:
@@ -661,14 +653,12 @@ class QueryDataParser:
                 node.debug_log()
         self.leaves = [leaf for node in self.nodes for leaf in node._iter_leaves()]
 
-        self.leaf_offset = 0
-
-    def parse(self, column: pa.Array):
+    def parse(self, column: pa.Array, state: QueryDataParserState):
         """Parse a single column response from VAST (see FieldNode.set for details)"""
-        if not self.leaf_offset < len(self.leaves):
-            raise ValueError(f'self.leaf_offset: {self.leaf_offset} are not < '
+        if not state.leaf_offset < len(self.leaves):
+            raise ValueError(f'state.leaf_offset: {state.leaf_offset} are not < '
                              f'than len(self.leaves): {len(self.leaves)}')
-        leaf = self.leaves[self.leaf_offset]
+        leaf = self.leaves[state.leaf_offset]
 
         # A column response may be sent in multiple chunks, therefore we need to combine
         # it into a single chunk to allow reconstruction using `Array.from_buffers()`.
@@ -685,13 +675,13 @@ class QueryDataParser:
             raise ValueError(f'len(array_list): {len(array_list)} are not eq '
                              f'with len(node_list): {len(node_list)}')
         for node, arr in zip(node_list, array_list):
-            node.set(arr)
+            node.set(arr, state)
 
-        self.leaf_offset += 1
+        state.leaf_offset += 1
 
-    def build(self) -> Optional[pa.Table]:
+    def build(self, state: QueryDataParserState) -> Optional[pa.Table]:
         """Try to build the resulting Table object (if all columns were parsed)"""
-        if self.leaf_offset < len(self.leaves):
+        if state.leaf_offset < len(self.leaves):
             return None
 
         if self.debug:
@@ -699,11 +689,12 @@ class QueryDataParser:
                 node.debug_log()
 
         result = pa.Table.from_arrays(
-            arrays=[node.build() for node in self.nodes],
+            arrays=[node.build(state) for node in self.nodes],
             schema=self.arrow_schema)
-        result.validate(full=self.debug) # does expensive validation checks only if debug is enabled
+        result.validate(full=self.debug)  # does expensive validation checks only if debug is enabled
         return result
 
+
 def _iter_nested_arrays(column: pa.Array) -> Iterator[pa.Array]:
     """Iterate over a single column response, and recursively generate all of its children."""
     yield column
@@ -715,7 +706,9 @@ def _iter_nested_arrays(column: pa.Array) -> Iterator[pa.Array]:
         yield from _iter_nested_arrays(column.values)  # Note: Map is serialized in VAST as a List<Struct<K, V>>
 
 
-TableInfo = namedtuple('table_info', 'name properties handle num_rows size_in_bytes')
+TableInfo = namedtuple('TableInfo', 'name properties handle num_rows size_in_bytes')
+
+
 def _parse_table_info(obj):
 
     name = obj.Name().decode()
@@ -725,6 +718,7 @@ def _parse_table_info(obj):
     used_bytes = obj.SizeInBytes()
     return TableInfo(name, properties, handle, num_rows, used_bytes)
 
+
 def build_record_batch(column_info, column_values):
     fields = [pa.field(column_name, column_type) for column_type, column_name in column_info]
     schema = pa.schema(fields)
@@ -732,6 +726,7 @@ def build_record_batch(column_info, column_values):
     batch = pa.record_batch(arrays, schema)
     return serialize_record_batch(batch)
 
+
 def serialize_record_batch(batch):
     sink = pa.BufferOutputStream()
     with pa.ipc.new_stream(sink, batch.schema) as writer:
@@ -739,61 +734,45 @@ def serialize_record_batch(batch):
     return sink.getvalue()
 
 # Results that returns from tablestats
-TableStatsResult = namedtuple("TableStatsResult",["num_rows", "size_in_bytes", "is_external_rowid_alloc", "endpoints"])
+
+
+TableStatsResult = namedtuple("TableStatsResult", ["num_rows", "size_in_bytes", "is_external_rowid_alloc", "endpoints"])
+
 
 class VastdbApi:
     # we expect the vast version to be <major>.<minor>.<patch>.<protocol>
     VAST_VERSION_REGEX = re.compile(r'^vast (\d+\.\d+\.\d+\.\d+)$')
 
-    def __init__(self, endpoint, access_key, secret_key, username=None, password=None,
-                 secure=False, auth_type=AuthType.SIGV4):
-        url_dict = urllib3.util.parse_url(endpoint)._asdict()
+    def __init__(self, endpoint, access_key, secret_key, auth_type=AuthType.SIGV4, ssl_verify=True):
+        url = urllib3.util.parse_url(endpoint)
         self.access_key = access_key
         self.secret_key = secret_key
-        self.username = username
-        self.password = password
-        self.secure = secure
-        self.auth_type = auth_type
-        self.executor_hosts = [endpoint]  # TODO: remove
-
-        username = username or ''
-        password = password or ''
-        if not url_dict['port']:
-            url_dict['port'] = 443 if secure else 80
-
-        self.port = url_dict['port']
 
         self.default_max_list_columns_page_size = 1000
         self.session = requests.Session()
-        self.session.verify = False
+        self.session.verify = ssl_verify
         self.session.headers['user-agent'] = "VastData Tabular API 1.0 - 2022 (c)"
-        if auth_type == AuthType.BASIC:
-            self.session.auth = requests.auth.HTTPBasicAuth(username, password)
-        else:
-            if url_dict['port'] != 80 and url_dict['port'] != 443:
-                self.aws_host = '{host}:{port}'.format(**url_dict)
-            else:
-                self.aws_host = '{host}'.format(**url_dict)
-
-            self.session.auth = AWSRequestsAuth(aws_access_key=access_key,
-                                                aws_secret_access_key=secret_key,
-                                                aws_host=self.aws_host,
-                                                aws_region='us-east-1',
-                                                aws_service='s3')
 
-        if not url_dict['scheme']:
-            url_dict['scheme'] = "https" if secure else "http"
+        if url.port in {80, 443, None}:
+            self.aws_host = f'{url.host}'
+        else:
+            self.aws_host = f'{url.host}:{url.port}'
 
-        url = urllib3.util.Url(**url_dict)
         self.url = str(url)
         _logger.debug('url=%s aws_host=%s', self.url, self.aws_host)
 
+        self.session.auth = AWSRequestsAuth(aws_access_key=access_key,
+                                            aws_secret_access_key=secret_key,
+                                            aws_host=self.aws_host,
+                                            aws_region='us-east-1',
+                                            aws_service='s3')
+
         # probe the cluster for its version
         self.vast_version = None
-        res = self.session.options(self.url)
+        res = self.session.get(self.url)
         server_header = res.headers.get("Server")
         if server_header is None:
-            _logger.error("OPTIONS response doesn't contain 'Server' header")
+            _logger.error("Response doesn't contain 'Server' header")
         else:
             _logger.debug("Server header is '%s'", server_header)
             if m := self.VAST_VERSION_REGEX.match(server_header):
@@ -994,9 +973,8 @@ class VastdbApi:
 
             return snapshots, is_truncated, marker
 
-
     def create_table(self, bucket, schema, name, arrow_schema, txid=0, client_tags=[], expected_retvals=[],
-                     topic_partitions=0, create_imports_table=False):
+                     topic_partitions=0, create_imports_table=False, use_external_row_ids_allocation=False):
         """
         Create a table, use the following request
         POST /bucket/schema/table?table HTTP/1.1
@@ -1017,6 +995,9 @@ class VastdbApi:
 
         serialized_schema = arrow_schema.serialize()
         headers['Content-Length'] = str(len(serialized_schema))
+        if use_external_row_ids_allocation:
+            headers['use-external-row-ids-alloc'] = str(use_external_row_ids_allocation)
+
         url_params = {'topic_partitions': str(topic_partitions)} if topic_partitions else {}
         if create_imports_table:
             url_params['sub-table'] = IMPORTED_OBJECTS_TABLE_NAME
@@ -1033,8 +1014,8 @@ class VastdbApi:
         if parquet_path:
             parquet_ds = pq.ParquetDataset(parquet_path)
         elif parquet_bucket_name and parquet_object_name:
-            s3fs  = pa.fs.S3FileSystem(access_key=self.access_key, secret_key=self.secret_key, endpoint_override=self.url)
-            parquet_ds = pq.ParquetDataset('/'.join([parquet_bucket_name,parquet_object_name]), filesystem=s3fs)
+            s3fs = pa.fs.S3FileSystem(access_key=self.access_key, secret_key=self.secret_key, endpoint_override=self.url)
+            parquet_ds = pq.ParquetDataset('/'.join([parquet_bucket_name, parquet_object_name]), filesystem=s3fs)
         else:
             raise RuntimeError(f'invalid params parquet_path={parquet_path} parquet_bucket_name={parquet_bucket_name} parquet_object_name={parquet_object_name}')
 
@@ -1049,7 +1030,6 @@ class VastdbApi:
         # create the table
         return self.create_table(bucket, schema, name, arrow_schema, txid, client_tags, expected_retvals)
 
-
     def get_table_stats(self, bucket, schema, name, txid=0, client_tags=[], expected_retvals=[]):
         """
         GET /mybucket/myschema/mytable?stats HTTP/1.1
@@ -1060,29 +1040,33 @@ class VastdbApi:
         """
         headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
         res = self.session.get(self._api_prefix(bucket=bucket, schema=schema, table=name, command="stats"), headers=headers)
-        if res.status_code == 200:
-            flatbuf = b''.join(res.iter_content(chunk_size=128))
-            stats = get_table_stats.GetRootAs(flatbuf)
-            num_rows = stats.NumRows()
-            size_in_bytes = stats.SizeInBytes()
-            is_external_rowid_alloc = stats.IsExternalRowidAlloc()
-            endpoints = []
-            if stats.VipsLength() == 0:
-                endpoints.append(self.url)
-            else:
-                ip_cls = IPv6Address if (stats.AddressType() == "ipv6") else IPv4Address
-                vips = [stats.Vips(i) for i in range(stats.VipsLength())]
-                ips = []
-                # extract the vips into list of IPs
-                for vip in vips:
-                    start_ip = int(ip_cls(vip.StartAddress().decode()))
-                    ips.extend(ip_cls(start_ip + i) for i  in range(vip.AddressCount()))
-                for ip in ips:
-                    prefix = "http" if not self.secure else "https"
-                    endpoints.append(f"{prefix}://{str(ip)}:{self.port}")
-            return TableStatsResult(num_rows, size_in_bytes, is_external_rowid_alloc, endpoints)
-
-        return self._check_res(res, "get_table_stats", expected_retvals)
+        self._check_res(res, "get_table_stats", expected_retvals)
+
+        flatbuf = b''.join(res.iter_content(chunk_size=128))
+        stats = get_table_stats.GetRootAs(flatbuf)
+        num_rows = stats.NumRows()
+        size_in_bytes = stats.SizeInBytes()
+        is_external_rowid_alloc = stats.IsExternalRowidAlloc()
+        endpoints = []
+        if stats.VipsLength() == 0:
+            endpoints.append(self.url)
+        else:
+            url = urllib3.util.parse_url(self.url)
+
+            ip_cls = IPv6Address if (stats.AddressType() == "ipv6") else IPv4Address
+            vips = [stats.Vips(i) for i in range(stats.VipsLength())]
+            ips = []
+            # extract the vips into list of IPs
+            for vip in vips:
+                start_ip = int(ip_cls(vip.StartAddress().decode()))
+                ips.extend(ip_cls(start_ip + i) for i in range(vip.AddressCount()))
+            # build a list of endpoint URLs, reusing schema and port (if specified when constructing the session).
+            # it is assumed that the client can access the returned IPs (e.g. if they are part of the VIP pool).
+            for ip in ips:
+                d = url._asdict()
+                d['host'] = str(ip)
+                endpoints.append(str(urllib3.util.Url(**d)))
+        return TableStatsResult(num_rows, size_in_bytes, is_external_rowid_alloc, tuple(endpoints))
 
     def alter_table(self, bucket, schema, name, txid=0, client_tags=[], table_properties="",
                     new_name="", expected_retvals=[]):
@@ -1171,7 +1155,6 @@ class VastdbApi:
 
             return bucket_name, schema_name, tables, next_key, is_truncated, count
 
-
     def add_columns(self, bucket, schema, name, arrow_schema, txid=0, client_tags=[], expected_retvals=[]):
         """
         Add a column to table, use the following request
@@ -1197,7 +1180,7 @@ class VastdbApi:
         return self._check_res(res, "add_columns", expected_retvals)
 
     def alter_column(self, bucket, schema, table, name, txid=0, client_tags=[], column_properties="",
-                     new_name="", column_sep = ".", column_stats="", expected_retvals=[]):
+                     new_name="", column_sep=".", column_stats="", expected_retvals=[]):
         """
         PUT /bucket/schema/table?column&tabular-column-name=ColumnName&tabular-new-column-name=NewColumnName HTTP/1.1
         Content-Length: ContentLength
@@ -1226,7 +1209,7 @@ class VastdbApi:
         headers['tabular-column-sep'] = column_sep
         headers['Content-Length'] = str(len(alter_column_req))
 
-        url_params = {'tabular-column-name': name }
+        url_params = {'tabular-column-name': name}
         if len(new_name):
             url_params['tabular-new-column-name'] = new_name
 
@@ -1573,7 +1556,7 @@ class VastdbApi:
         return self._check_res(res, "import_data", expected_retvals)
 
     def _record_batch_slices(self, batch, rows_per_slice=None):
-        max_slice_size_in_bytes = int(0.9*5*1024*1024) # 0.9 * 5MB
+        max_slice_size_in_bytes = int(0.9 * 5 * 1024 * 1024)  # 0.9 * 5MB
         batch_len = len(batch)
         serialized_batch = serialize_record_batch(batch)
         batch_size_in_bytes = len(serialized_batch)
@@ -1591,10 +1574,10 @@ class VastdbApi:
             # Attempt slicing according to the current rows_per_slice
             offset = 0
             serialized_slices = []
-            for i in range(math.ceil(batch_len/rows_per_slice)):
+            for i in range(math.ceil(batch_len / rows_per_slice)):
                 offset = rows_per_slice * i
                 if offset >= batch_len:
-                    done_slicing=True
+                    done_slicing = True
                     break
                 slice_batch = batch.slice(offset, rows_per_slice)
                 serialized_slice_batch = serialize_record_batch(slice_batch)
@@ -1605,7 +1588,7 @@ class VastdbApi:
                 else:
                     _logger.info(f'Using rows_per_slice {rows_per_slice} slice {i} size {sizeof_serialized_slice_batch} exceeds {max_slice_size_in_bytes} bytes, trying smaller rows_per_slice')
                     # We have a slice that is too large
-                    rows_per_slice = int(rows_per_slice/2)
+                    rows_per_slice = int(rows_per_slice / 2)
                     if rows_per_slice < 1:
                         raise ValueError('cannot decrease batch size below 1 row')
                     break
@@ -1628,7 +1611,8 @@ class VastdbApi:
         headers['Content-Length'] = str(len(record_batch))
         res = self.session.post(self._api_prefix(bucket=bucket, schema=schema, table=table, command="rows"),
                                 data=record_batch, headers=headers, stream=True)
-        return self._check_res(res, "insert_rows", expected_retvals)
+        self._check_res(res, "insert_rows", expected_retvals)
+        res.raw.read()  # flush the response
 
     def update_rows(self, bucket, schema, table, record_batch, txid=0, client_tags=[], expected_retvals=[]):
         """
@@ -1644,7 +1628,7 @@ class VastdbApi:
         headers['Content-Length'] = str(len(record_batch))
         res = self.session.put(self._api_prefix(bucket=bucket, schema=schema, table=table, command="rows"),
                                 data=record_batch, headers=headers)
-        return self._check_res(res, "update_rows", expected_retvals)
+        self._check_res(res, "update_rows", expected_retvals)
 
     def delete_rows(self, bucket, schema, table, record_batch, txid=0, client_tags=[], expected_retvals=[],
                     delete_from_imports_table=False):
@@ -1663,7 +1647,7 @@ class VastdbApi:
 
         res = self.session.delete(self._api_prefix(bucket=bucket, schema=schema, table=table, command="rows", url_params=url_params),
                                   data=record_batch, headers=headers)
-        return self._check_res(res, "delete_rows", expected_retvals)
+        self._check_res(res, "delete_rows", expected_retvals)
 
     def create_projection(self, bucket, schema, table, name, columns, txid=0, client_tags=[], expected_retvals=[]):
         """
@@ -1873,6 +1857,10 @@ class VastdbApi:
             return columns, next_key, is_truncated, count
 
 
+class QueryDataInternalError(Exception):
+    pass
+
+
 def _iter_query_data_response_columns(fileobj, stream_ids=None):
     readers = {}  # {stream_id: pa.ipc.RecordBatchStreamReader}
     while True:
@@ -1897,8 +1885,8 @@ def _iter_query_data_response_columns(fileobj, stream_ids=None):
         if stream_id == TABULAR_QUERY_DATA_FAILED_STREAM_ID:
             # read the terminating end chunk from socket
             res = fileobj.read()
-            _logger.warning("stream_id=%d res=%s (failed)", stream_id, res)
-            raise IOError(f"Query data stream failed res={res}")
+            _logger.debug("stream_id=%d res=%s (failed)", stream_id, res)
+            raise QueryDataInternalError()  # connection closed by server due to an internal error
 
         next_row_id_bytes = fileobj.read(8)
         next_row_id, = struct.unpack('<Q', next_row_id_bytes)
@@ -1913,7 +1901,7 @@ def _iter_query_data_response_columns(fileobj, stream_ids=None):
 
         (reader, batches) = readers[stream_id]
         try:
-            batch = reader.read_next_batch() # read single-column chunk data
+            batch = reader.read_next_batch()  # read single-column chunk data
             _logger.debug("stream_id=%d rows=%d chunk=%s", stream_id, len(batch), batch)
             batches.append(batch)
         except StopIteration:  # we got an end-of-stream IPC message for a given stream ID
@@ -1923,7 +1911,7 @@ def _iter_query_data_response_columns(fileobj, stream_ids=None):
             yield (stream_id, next_row_id, table)
 
 
-def parse_query_data_response(conn, schema, stream_ids=None, start_row_ids=None, debug=False):
+def parse_query_data_response(conn, schema, stream_ids=None, start_row_ids=None, debug=False, parser: Optional[QueryDataParser] = None):
     """
     Generates pyarrow.Table objects from QueryData API response stream.
 
@@ -1933,16 +1921,18 @@ def parse_query_data_response(conn, schema, stream_ids=None, start_row_ids=None,
         start_row_ids = {}
 
     is_empty_projection = (len(schema) == 0)
-    parsers = defaultdict(lambda: QueryDataParser(schema, debug=debug))  # {stream_id: QueryDataParser}
+    if parser is None:
+        parser = QueryDataParser(schema, debug=debug)
+    states: Dict[int, QueryDataParser.QueryDataParserState] = defaultdict(lambda: QueryDataParser.QueryDataParserState())  # {stream_id: QueryDataParser}
 
     for stream_id, next_row_id, table in _iter_query_data_response_columns(conn, stream_ids):
-        parser = parsers[stream_id]
+        state = states[stream_id]
         for column in table.columns:
-            parser.parse(column)
+            parser.parse(column, state)
 
-        parsed_table = parser.build()
+        parsed_table = parser.build(state)
         if parsed_table is not None:  # when we got all columns (and before starting a new "select_rows" cycle)
-            parsers.pop(stream_id)
+            states.pop(stream_id)
             if is_empty_projection:  # VAST returns an empty RecordBatch, with the correct rows' count
                 parsed_table = table
 
@@ -1951,8 +1941,9 @@ def parse_query_data_response(conn, schema, stream_ids=None, start_row_ids=None,
             start_row_ids[stream_id] = next_row_id
             yield parsed_table  # the result of a single "select_rows()" cycle
 
-    if parsers:
-        raise EOFError(f'all streams should be done before EOF. {parsers}')
+    if states:
+        raise EOFError(f'all streams should be done before EOF. {states}')
+
 
 def get_field_type(builder: flatbuffers.Builder, field: pa.Field):
     if field.type.equals(pa.int64()):
@@ -2095,6 +2086,7 @@ def get_field_type(builder: flatbuffers.Builder, field: pa.Field):
 
     return field_type, field_type_type
 
+
 def build_field(builder: flatbuffers.Builder, f: pa.Field, name: str):
     children = None
     if isinstance(f.type, pa.StructType):
@@ -2142,12 +2134,13 @@ def build_field(builder: flatbuffers.Builder, f: pa.Field, name: str):
 
 
 class QueryDataRequest:
-    def __init__(self, serialized, response_schema):
+    def __init__(self, serialized, response_schema, response_parser):
         self.serialized = serialized
         self.response_schema = response_schema
+        self.response_parser = response_parser
 
 
-def build_query_data_request(schema: 'pa.Schema' = pa.schema([]), predicate: ibis.expr.types.BooleanColumn = None, field_names: list = None):
+def build_query_data_request(schema: 'pa.Schema' = pa.schema([]), predicate: ibis.expr.types.BooleanColumn = None, field_names: Optional[List[str]] = None):
     builder = flatbuffers.Builder(1024)
 
     source_name = builder.CreateString('')  # required
@@ -2201,7 +2194,8 @@ def build_query_data_request(schema: 'pa.Schema' = pa.schema([]), predicate: ibi
     relation = fb_relation.End(builder)
 
     builder.Finish(relation)
-    return QueryDataRequest(serialized=builder.Output(), response_schema=response_schema)
+
+    return QueryDataRequest(serialized=builder.Output(), response_schema=response_schema, response_parser=QueryDataParser(response_schema))
 
 
 def convert_column_types(table: 'pa.Table') -> 'pa.Table':