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

vastdb/internal_commands.py

@@ -1,14 +1,13 @@
 import itertools
 import json
 import logging
-import math
 import re
 import struct
 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 +91,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 +126,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 +156,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,12 +171,17 @@ 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,
             Greater,
             GreaterEqual,
+            InValues,
             Less,
             LessEqual,
             Not,
@@ -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,40 +219,54 @@ 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))
+                    op_type = type(inner_op)
+                    builder_func: Any = builder_map.get(op_type)
                     if not builder_func:
-                        raise NotImplementedError(inner_op.name)
+                        if op_type == InValues:
+                            builder_func = self.build_equal
+                        else:
+                            raise NotImplementedError(self.expr)
 
                     if builder_func == self.build_is_null:
                         column, = inner_op.args
-                        literal = None
+                        literals = (None,)
                     elif builder_func == self.build_is_not_null:
                         not_arg, = inner_op.args
                         # currently we only support not is_null, checking we really got is_null under the not:
                         if not builder_map.get(type(not_arg)) == self.build_is_null:
-                            raise NotImplementedError(not_arg.args[0].name)
+                            raise NotImplementedError(self.expr)
                         column, = not_arg.args
-                        literal = None
+                        literals = (None,)
                     else:
-                        column, literal = inner_op.args
-                        if not isinstance(literal, Literal):
-                            raise NotImplementedError(inner_op.name)
+                        column, arg = inner_op.args
+                        if isinstance(arg, tuple):
+                            literals = arg
+                        else:
+                            literals = (arg,)
+                        for literal in literals:
+                            if not isinstance(literal, Literal):
+                                raise NotImplementedError(self.expr)
 
                     if not isinstance(column, TableColumn):
-                        raise NotImplementedError(inner_op.name)
+                        raise NotImplementedError(self.expr)
 
                     field_name = column.name
                     if prev_field_name is None:
                         prev_field_name = field_name
                     elif prev_field_name != field_name:
-                        raise NotImplementedError(op.name)
+                        raise NotImplementedError(self.expr)
 
-                    args_offsets = [self.build_column(position=positions_map[field_name])]
-                    if literal:
-                        field = self.schema.field(field_name)
-                        args_offsets.append(self.build_literal(field=field, value=literal.value))
+                    column_offset = self.build_column(position=positions_map[field_name])
+                    field = self.schema.field(field_name)
+                    for literal in literals:
+                        args_offsets = [column_offset]
+                        if literal is not None:
+                            args_offsets.append(self.build_literal(field=field, value=literal.value))
 
-                    inner_offsets.append(builder_func(*args_offsets))
+                        inner_offsets.append(builder_func(*args_offsets))
+
+                if not inner_offsets:
+                    raise NotImplementedError(self.expr)  # an empty OR is equivalent to a 'FALSE' literal
 
                 domain_offset = self.build_or(inner_offsets)
                 offsets.append(domain_offset)
@@ -270,20 +288,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 +343,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 +557,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 +587,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 +608,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 +627,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 +667,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 +689,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 +703,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 +720,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,75 +732,47 @@ 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)
-    arrays = [pa.array(column_values[column_type], type=column_type) for column_type, _ in column_info]
-    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:
-        writer.write(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):
@@ -973,30 +952,30 @@ class VastdbApi:
 
             return bucket_name, schemas, next_key, is_truncated, count
 
-    def list_snapshots(self, bucket, max_keys=1000, next_token=None, expected_retvals=None):
+    def list_snapshots(self, bucket, max_keys=1000, next_token=None, name_prefix=''):
         next_token = next_token or ''
-        expected_retvals = expected_retvals or []
-        url_params = {'list_type': '2', 'prefix': '.snapshot/', 'delimiter': '/', 'max_keys': str(max_keys)}
+        url_params = {'list_type': '2', 'prefix': '.snapshot/' + name_prefix, 'delimiter': '/', 'max_keys': str(max_keys)}
         if next_token:
             url_params['continuation-token'] = next_token
 
         res = self.session.get(self._api_prefix(bucket=bucket, command="list", url_params=url_params), headers={}, stream=True)
-        self._check_res(res, "list_snapshots", expected_retvals)
-        if res.status_code == 200:
-            out = b''.join(res.iter_content(chunk_size=128))
-            xml_str = out.decode()
-            xml_dict = xmltodict.parse(xml_str)
-            list_res = xml_dict['ListBucketResult']
-            is_truncated = list_res['IsTruncated'] == 'true'
-            marker = list_res['Marker']
-            common_prefixes = list_res['CommonPrefixes'] if 'CommonPrefixes' in list_res else []
-            snapshots = [v['Prefix'] for v in common_prefixes]
+        self._check_res(res, "list_snapshots")
 
-            return snapshots, is_truncated, marker
+        out = b''.join(res.iter_content(chunk_size=128))
+        xml_str = out.decode()
+        xml_dict = xmltodict.parse(xml_str)
+        list_res = xml_dict['ListBucketResult']
+        is_truncated = list_res['IsTruncated'] == 'true'
+        marker = list_res['Marker']
+        common_prefixes = list_res.get('CommonPrefixes', [])
+        if isinstance(common_prefixes, dict):  # in case there is a single snapshot
+            common_prefixes = [common_prefixes]
+        snapshots = [v['Prefix'] for v in common_prefixes]
 
+        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 +996,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 +1015,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,8 +1031,7 @@ 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=[]):
+    def get_table_stats(self, bucket, schema, name, txid=0, client_tags=[], expected_retvals=[], imports_table_stats=False):
         """
         GET /mybucket/myschema/mytable?stats HTTP/1.1
         tabular-txid: TransactionId
@@ -1059,30 +1040,35 @@ class VastdbApi:
         The Command will return the statistics in flatbuf format
         """
         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)
+        url_params = {'sub-table': IMPORTED_OBJECTS_TABLE_NAME} if imports_table_stats else {}
+        res = self.session.get(self._api_prefix(bucket=bucket, schema=schema, table=name, command="stats", url_params=url_params), headers=headers)
+        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 +1157,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 +1182,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 +1211,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
 
@@ -1544,11 +1529,18 @@ class VastdbApi:
             if response.status_code != 200:
                 return response
 
+            ALLOWED_IMPORT_STATES = {
+                'Success',
+                'TabularInProgress',
+                'TabularAlreadyImported',
+                'TabularImportNotStarted',
+            }
+
             chunk_size = 1024
             for chunk in response.iter_content(chunk_size=chunk_size):
                 chunk_dict = json.loads(chunk)
                 _logger.debug("import data chunk=%s, result: %s", chunk_dict, chunk_dict['res'])
-                if chunk_dict['res'] != 'Success' and chunk_dict['res'] != 'TabularInProgress' and chunk_dict['res'] != 'TabularAlreadyImported':
+                if chunk_dict['res'] not in ALLOWED_IMPORT_STATES:
                     raise errors.ImportFilesError(
                         f"Encountered an error during import_data. status: {chunk_dict['res']}, "
                         f"error message: {chunk_dict['err_msg'] or 'Unexpected error'} during import of "
@@ -1572,48 +1564,6 @@ 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
-        batch_len = len(batch)
-        serialized_batch = serialize_record_batch(batch)
-        batch_size_in_bytes = len(serialized_batch)
-        _logger.debug('max_slice_size_in_bytes=%d batch_len=%d batch_size_in_bytes=%d',
-                      max_slice_size_in_bytes, batch_len, batch_size_in_bytes)
-
-        if not rows_per_slice:
-            if batch_size_in_bytes < max_slice_size_in_bytes:
-                rows_per_slice = batch_len
-            else:
-                rows_per_slice = int(0.9 * batch_len * max_slice_size_in_bytes / batch_size_in_bytes)
-
-        done_slicing = False
-        while not done_slicing:
-            # Attempt slicing according to the current rows_per_slice
-            offset = 0
-            serialized_slices = []
-            for i in range(math.ceil(batch_len/rows_per_slice)):
-                offset = rows_per_slice * i
-                if offset >= batch_len:
-                    done_slicing=True
-                    break
-                slice_batch = batch.slice(offset, rows_per_slice)
-                serialized_slice_batch = serialize_record_batch(slice_batch)
-                sizeof_serialized_slice_batch = len(serialized_slice_batch)
-
-                if sizeof_serialized_slice_batch <= max_slice_size_in_bytes:
-                    serialized_slices.append(serialized_slice_batch)
-                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)
-                    if rows_per_slice < 1:
-                        raise ValueError('cannot decrease batch size below 1 row')
-                    break
-            else:
-                done_slicing = True
-
-        return serialized_slices
-
     def insert_rows(self, bucket, schema, table, record_batch, txid=0, client_tags=[], expected_retvals=[]):
         """
         POST /mybucket/myschema/mytable?rows HTTP/1.1
@@ -1628,7 +1578,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 +1595,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 +1614,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 +1824,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 +1852,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 +1868,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 +1878,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 +1888,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 +1908,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 +2053,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 +2101,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 +2161,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':