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

vastdb/internal_commands.py

@@ -1,26 +1,23 @@
+import itertools
+import json
 import logging
+import math
+import re
 import struct
 import urllib.parse
 from collections import defaultdict, namedtuple
-from datetime import datetime
 from enum import Enum
-from typing import Union, Optional, Iterator
-import ibis
-import xmltodict
-import math
-from functools import cmp_to_key
-import pyarrow.parquet as pq
+from ipaddress import IPv4Address, IPv6Address
+from typing import Any, Dict, Iterator, List, Optional, Union
+
 import flatbuffers
+import ibis
 import pyarrow as pa
+import pyarrow.parquet as pq
 import requests
-import json
-import itertools
-from aws_requests_auth.aws_auth import AWSRequestsAuth
 import urllib3
-import re
-
-from . import errors
-from ipaddress import IPv4Address, IPv6Address
+import xmltodict
+from aws_requests_auth.aws_auth import AWSRequestsAuth
 
 import vast_flatbuf.org.apache.arrow.computeir.flatbuf.BinaryLiteral as fb_binary_lit
 import vast_flatbuf.org.apache.arrow.computeir.flatbuf.BooleanLiteral as fb_bool_lit
@@ -32,10 +29,10 @@ import vast_flatbuf.org.apache.arrow.computeir.flatbuf.FieldIndex as fb_field_in
 import vast_flatbuf.org.apache.arrow.computeir.flatbuf.FieldRef as fb_field_ref
 import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Float32Literal as fb_float32_lit
 import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Float64Literal as fb_float64_lit
+import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Int8Literal as fb_int8_lit
 import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Int16Literal as fb_int16_lit
 import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Int32Literal as fb_int32_lit
 import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Int64Literal as fb_int64_lit
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Int8Literal as fb_int8_lit
 import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Literal as fb_literal
 import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Relation as fb_relation
 import vast_flatbuf.org.apache.arrow.computeir.flatbuf.RelationImpl as rel_impl
@@ -48,45 +45,54 @@ import vast_flatbuf.org.apache.arrow.flatbuf.Bool as fb_bool
 import vast_flatbuf.org.apache.arrow.flatbuf.Date as fb_date
 import vast_flatbuf.org.apache.arrow.flatbuf.Decimal as fb_decimal
 import vast_flatbuf.org.apache.arrow.flatbuf.Field as fb_field
+import vast_flatbuf.org.apache.arrow.flatbuf.FixedSizeBinary as fb_fixed_size_binary
 import vast_flatbuf.org.apache.arrow.flatbuf.FloatingPoint as fb_floating_point
 import vast_flatbuf.org.apache.arrow.flatbuf.Int as fb_int
-import vast_flatbuf.org.apache.arrow.flatbuf.Schema as fb_schema
-import vast_flatbuf.org.apache.arrow.flatbuf.Time as fb_time
-import vast_flatbuf.org.apache.arrow.flatbuf.Struct_ as fb_struct
 import vast_flatbuf.org.apache.arrow.flatbuf.List as fb_list
 import vast_flatbuf.org.apache.arrow.flatbuf.Map as fb_map
-import vast_flatbuf.org.apache.arrow.flatbuf.FixedSizeBinary as fb_fixed_size_binary
+import vast_flatbuf.org.apache.arrow.flatbuf.Schema as fb_schema
+import vast_flatbuf.org.apache.arrow.flatbuf.Struct_ as fb_struct
+import vast_flatbuf.org.apache.arrow.flatbuf.Time as fb_time
 import vast_flatbuf.org.apache.arrow.flatbuf.Timestamp as fb_timestamp
 import vast_flatbuf.org.apache.arrow.flatbuf.Utf8 as fb_utf8
 import vast_flatbuf.tabular.AlterColumnRequest as tabular_alter_column
+import vast_flatbuf.tabular.AlterProjectionTableRequest as tabular_alter_projection
 import vast_flatbuf.tabular.AlterSchemaRequest as tabular_alter_schema
 import vast_flatbuf.tabular.AlterTableRequest as tabular_alter_table
-import vast_flatbuf.tabular.AlterProjectionTableRequest as tabular_alter_projection
+import vast_flatbuf.tabular.Column as tabular_projecion_column
+import vast_flatbuf.tabular.ColumnType as tabular_proj_column_type
+import vast_flatbuf.tabular.CreateProjectionRequest as tabular_create_projection
 import vast_flatbuf.tabular.CreateSchemaRequest as tabular_create_schema
 import vast_flatbuf.tabular.ImportDataRequest as tabular_import_data
 import vast_flatbuf.tabular.S3File as tabular_s3_file
-import vast_flatbuf.tabular.CreateProjectionRequest as tabular_create_projection
-import vast_flatbuf.tabular.Column as tabular_projecion_column
-import vast_flatbuf.tabular.ColumnType as tabular_proj_column_type
-
 from vast_flatbuf.org.apache.arrow.computeir.flatbuf.Deref import Deref
-from vast_flatbuf.org.apache.arrow.computeir.flatbuf.ExpressionImpl import ExpressionImpl
+from vast_flatbuf.org.apache.arrow.computeir.flatbuf.ExpressionImpl import (
+    ExpressionImpl,
+)
 from vast_flatbuf.org.apache.arrow.computeir.flatbuf.LiteralImpl import LiteralImpl
 from vast_flatbuf.org.apache.arrow.flatbuf.DateUnit import DateUnit
 from vast_flatbuf.org.apache.arrow.flatbuf.TimeUnit import TimeUnit
 from vast_flatbuf.org.apache.arrow.flatbuf.Type import Type
+from vast_flatbuf.tabular.GetProjectionTableStatsResponse import (
+    GetProjectionTableStatsResponse as get_projection_table_stats,
+)
+from vast_flatbuf.tabular.GetTableStatsResponse import (
+    GetTableStatsResponse as get_table_stats,
+)
+from vast_flatbuf.tabular.ListProjectionsResponse import (
+    ListProjectionsResponse as list_projections,
+)
 from vast_flatbuf.tabular.ListSchemasResponse import ListSchemasResponse as list_schemas
 from vast_flatbuf.tabular.ListTablesResponse import ListTablesResponse as list_tables
-from vast_flatbuf.tabular.GetTableStatsResponse import GetTableStatsResponse as get_table_stats
-from vast_flatbuf.tabular.GetProjectionTableStatsResponse import GetProjectionTableStatsResponse as get_projection_table_stats
-from vast_flatbuf.tabular.ListProjectionsResponse import ListProjectionsResponse as list_projections
+
+from . import errors
 
 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"
 
@@ -121,18 +127,11 @@ def get_unit_to_flatbuff_time_unit(type):
     }
     return unit_to_flatbuff_time_unit[type]
 
-class Predicate:
-    unit_to_epoch = {
-        'ns': 1_000_000,
-        'us': 1_000,
-        'ms': 1,
-        's': 0.001
-    }
 
+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)
@@ -158,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:
@@ -173,8 +172,22 @@ class Predicate:
         return builder.EndVector()
 
     def serialize(self, builder: 'flatbuffers.builder.Builder'):
-        from ibis.expr.operations.generic import TableColumn, Literal, IsNull
-        from ibis.expr.operations.logical import Greater, GreaterEqual, Less, LessEqual, Equals, NotEquals, And, Or, Not
+        from ibis.expr.operations.generic import (
+            IsNull,
+            Literal,
+            TableColumn,
+        )
+        from ibis.expr.operations.logical import (
+            And,
+            Equals,
+            Greater,
+            GreaterEqual,
+            Less,
+            LessEqual,
+            Not,
+            NotEquals,
+            Or,
+        )
         from ibis.expr.operations.strings import StringContains
 
         builder_map = {
@@ -189,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
 
@@ -206,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)
 
@@ -261,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,
@@ -330,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
@@ -403,7 +404,7 @@ class Predicate:
             field_type = fb_utf8.End(self.builder)
 
             value = self.builder.CreateString(value)
-        elif field.type.equals(pa.date32()):  # pa.date64()
+        elif field.type.equals(pa.date32()):  # pa.date64() is not supported
             literal_type = fb_date32_lit
             literal_impl = LiteralImpl.DateLiteral
 
@@ -411,37 +412,49 @@ class Predicate:
             fb_date.Start(self.builder)
             fb_date.AddUnit(self.builder, DateUnit.DAY)
             field_type = fb_date.End(self.builder)
-
-            start_date = datetime.fromtimestamp(0).date()
-            date_delta = value - start_date
-            value = date_delta.days
+            value, = pa.array([value], field.type).cast(pa.int32()).to_pylist()
         elif isinstance(field.type, pa.TimestampType):
             literal_type = fb_timestamp_lit
             literal_impl = LiteralImpl.TimestampLiteral
 
+            if field.type.equals(pa.timestamp('s')):
+                unit = TimeUnit.SECOND
+            if field.type.equals(pa.timestamp('ms')):
+                unit = TimeUnit.MILLISECOND
+            if field.type.equals(pa.timestamp('us')):
+                unit = TimeUnit.MICROSECOND
+            if field.type.equals(pa.timestamp('ns')):
+                unit = TimeUnit.NANOSECOND
+
             field_type_type = Type.Timestamp
             fb_timestamp.Start(self.builder)
-            fb_timestamp.AddUnit(self.builder, get_unit_to_flatbuff_time_unit(field.type.unit))
+            fb_timestamp.AddUnit(self.builder, unit)
             field_type = fb_timestamp.End(self.builder)
-
-            value = int(int(value) * self.unit_to_epoch[field.type.unit])
-        elif field.type.equals(pa.time32('s')) or field.type.equals(pa.time32('ms')) or field.type.equals(pa.time64('us')) or field.type.equals(pa.time64('ns')):
-
+            value, = pa.array([value], field.type).cast(pa.int64()).to_pylist()
+        elif isinstance(field.type, (pa.Time32Type, pa.Time64Type)):
             literal_type = fb_time_lit
             literal_impl = LiteralImpl.TimeLiteral
 
-            field_type_str = str(field.type)
-            start = field_type_str.index('[')
-            end = field_type_str.index(']')
-            unit = field_type_str[start + 1:end]
+            if field.type.equals(pa.time32('s')):
+                target_type = pa.int32()
+                unit = TimeUnit.SECOND
+            if field.type.equals(pa.time32('ms')):
+                target_type = pa.int32()
+                unit = TimeUnit.MILLISECOND
+            if field.type.equals(pa.time64('us')):
+                target_type = pa.int64()
+                unit = TimeUnit.MICROSECOND
+            if field.type.equals(pa.time64('ns')):
+                target_type = pa.int64()
+                unit = TimeUnit.NANOSECOND
 
             field_type_type = Type.Time
             fb_time.Start(self.builder)
             fb_time.AddBitWidth(self.builder, field.type.bit_width)
-            fb_time.AddUnit(self.builder, get_unit_to_flatbuff_time_unit(unit))
+            fb_time.AddUnit(self.builder, unit)
             field_type = fb_time.End(self.builder)
 
-            value = int(value) * self.unit_to_epoch[unit]
+            value, = pa.array([value], field.type).cast(target_type).to_pylist()
         elif field.type.equals(pa.bool_()):
             literal_type = fb_bool_lit
             literal_impl = LiteralImpl.BooleanLiteral
@@ -530,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]
@@ -553,13 +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.is_projected = False
-        self.projected_field = self.field
+            self.children = []  # for non-nested types
 
     def _iter_to_root(self) -> Iterator['FieldNode']:
         yield self
@@ -580,24 +594,14 @@ class FieldNode:
             for child in self.children:
                 yield from child._iter_leaves()
 
-    def _iter_projected_leaves(self) -> Iterator['FieldNode']:
-        """Generate only leaf nodes (i.e. columns having scalar types)."""
-        if not self.children:
-            if self.is_projected:
-                yield self
-        else:
-            for child in self.children:
-                if child.is_projected:
-                    yield from child._iter_projected_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.
 
@@ -609,68 +613,52 @@ 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 if node.is_projected]
-        _logger.debug('build: self.field.name=%s, self.projected_field.type=%s, self.length=%s, self.buffers=%s children=%s',
-                      self.field.name, self.projected_field.type, self.length, self.buffers, children)
-        result = pa.Array.from_buffers(self.projected_field.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
 
-    def build_projected_field(self):
-        if isinstance(self.type, pa.StructType):
-            [child.build_projected_field() for child in self.children if child.is_projected]
-            self.projected_field = pa.field(self.field.name,
-                                            pa.struct([child.projected_field for child in self.children if child.is_projected]),
-                                            self.field.nullable,
-                                            self.field.metadata)
 
 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, projection_positions=None):
+    def __init__(self, arrow_schema: pa.Schema, *, debug=False):
         self.arrow_schema = arrow_schema
-        self.projection_positions = projection_positions
-        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:
             for node in self.nodes:
                 node.debug_log()
         self.leaves = [leaf for node in self.nodes for leaf in node._iter_leaves()]
-        self.mark_projected_nodes()
-        [node.build_projected_field() for node in self.nodes]
-        self.projected_leaves = [leaf for node in self.nodes for leaf in node._iter_projected_leaves()]
-
-        self.leaf_offset = 0
 
-    def mark_projected_nodes(self):
-        for leaf in self.leaves:
-            if self.projection_positions is None or leaf.index in self.projection_positions:
-                for node in leaf._iter_to_root():
-                    node.is_projected = True
-
-    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.projected_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.projected_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()`.
@@ -687,38 +675,26 @@ 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, output_field_names=None) -> 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.projection_positions is not None:
-            if self.leaf_offset < len(self.projection_positions):
-                return None
-        else:
-            if self.leaf_offset < len(self.leaves):
-                return None
+        if state.leaf_offset < len(self.leaves):
+            return None
 
         if self.debug:
             for node in self.nodes:
                 node.debug_log()
 
-        # sort resulting table according to the output field names
-        projected_nodes = [node for node in self.nodes if node.is_projected]
-        if output_field_names is not None:
-            def key_func(projected_node):
-                return output_field_names.index(projected_node.field.name)
-            sorted_projected_nodes = sorted(projected_nodes, key=key_func)
-        else:
-            sorted_projected_nodes = projected_nodes
-
         result = pa.Table.from_arrays(
-            arrays=[node.build() for node in sorted_projected_nodes],
-            schema = pa.schema([node.projected_field for node in sorted_projected_nodes]))
-        result.validate(full=True) # does expensive validation checks only if debug is enabled
+            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
         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
@@ -730,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()
@@ -740,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)
@@ -747,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:
@@ -754,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):
@@ -1009,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
@@ -1032,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
@@ -1048,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}')
 
@@ -1064,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
@@ -1075,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=[]):
@@ -1186,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
@@ -1212,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
@@ -1241,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
 
@@ -1588,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)
@@ -1606,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)
@@ -1620,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
@@ -1643,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=[]):
         """
@@ -1659,9 +1628,10 @@ 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=[]):
+    def delete_rows(self, bucket, schema, table, record_batch, txid=0, client_tags=[], expected_retvals=[],
+                    delete_from_imports_table=False):
         """
         DELETE /mybucket/myschema/mytable?rows HTTP/1.1
         Content-Length: ContentLength
@@ -1673,9 +1643,11 @@ class VastdbApi:
         """
         headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
         headers['Content-Length'] = str(len(record_batch))
-        res = self.session.delete(self._api_prefix(bucket=bucket, schema=schema, table=table, command="rows"),
-                               data=record_batch, headers=headers)
-        return self._check_res(res, "delete_rows", expected_retvals)
+        url_params = {'sub-table': IMPORTED_OBJECTS_TABLE_NAME} if delete_from_imports_table else {}
+
+        res = self.session.delete(self._api_prefix(bucket=bucket, schema=schema, table=table, command="rows", url_params=url_params),
+                                  data=record_batch, headers=headers)
+        self._check_res(res, "delete_rows", expected_retvals)
 
     def create_projection(self, bucket, schema, table, name, columns, txid=0, client_tags=[], expected_retvals=[]):
         """
@@ -1885,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:
@@ -1909,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)
@@ -1925,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
@@ -1935,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.
 
@@ -1943,20 +1919,20 @@ def parse_query_data_response(conn, schema, stream_ids=None, start_row_ids=None,
     """
     if start_row_ids is None:
         start_row_ids = {}
-    projection_positions = schema.projection_positions
-    arrow_schema = schema.arrow_schema
-    output_field_names = schema.output_field_names
-    _logger.debug(f'projection_positions={projection_positions} len(arrow_schema)={len(arrow_schema)} arrow_schema={arrow_schema}')
-    is_empty_projection = (len(projection_positions) == 0)
-    parsers = defaultdict(lambda: QueryDataParser(arrow_schema, debug=debug, projection_positions=projection_positions))  # {stream_id: QueryDataParser}
+
+    is_empty_projection = (len(schema) == 0)
+    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(output_field_names)
+        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
 
@@ -1965,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()):
@@ -2042,7 +2019,7 @@ def get_field_type(builder: flatbuffers.Builder, field: pa.Field):
         fb_utf8.Start(builder)
         field_type = fb_utf8.End(builder)
 
-    elif field.type.equals(pa.date32()):  # pa.date64()
+    elif field.type.equals(pa.date32()):  # pa.date64() is not supported
         field_type_type = Type.Date
         fb_date.Start(builder)
         fb_date.AddUnit(builder, DateUnit.DAY)
@@ -2109,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):
@@ -2155,19 +2133,14 @@ def build_field(builder: flatbuffers.Builder, f: pa.Field, name: str):
     return fb_field.End(builder)
 
 
-class VastDBResponseSchema:
-    def __init__(self, arrow_schema, projection_positions, output_field_names):
-        self.arrow_schema = arrow_schema
-        self.projection_positions = projection_positions
-        self.output_field_names = output_field_names
-
 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
@@ -2187,31 +2160,17 @@ def build_query_data_request(schema: 'pa.Schema' = pa.schema([]), predicate: ibi
     filter_obj = predicate.serialize(builder)
 
     parser = QueryDataParser(schema)
-    leaves_map = {}
-    for node in parser.nodes:
-        for descendent in node._iter_nodes():
-            if descendent.parent and isinstance(descendent.parent.type, (pa.ListType, pa.MapType)):
-                continue
-            iter_from_root = reversed(list(descendent._iter_to_root()))
-            descendent_full_name = '.'.join([n.field.name for n in iter_from_root])
-            descendent_leaves = [leaf.index for leaf in descendent._iter_leaves()]
-            leaves_map[descendent_full_name] = descendent_leaves
-
-    output_field_names = None
+    fields_map = {node.field.name: node.field for node in parser.nodes}
+    leaves_map = {node.field.name: [leaf.index for leaf in node._iter_leaves()] for node in parser.nodes}
+
     if field_names is None:
         field_names = [field.name for field in schema]
-    else:
-        output_field_names  = [f.split('.')[0] for f in field_names]
-        # sort projected field_names according to positions to maintain ordering according to the schema
-        def compare_field_names_by_pos(field_name1, field_name2):
-            return leaves_map[field_name1][0]-leaves_map[field_name2][0]
-        field_names = sorted(field_names, key=cmp_to_key(compare_field_names_by_pos))
 
+    response_schema = pa.schema([fields_map[name] for name in field_names])
     projection_fields = []
-    projection_positions = []
     for field_name in field_names:
+        # TODO: only root-level projection pushdown is supported (i.e. no support for SELECT s.x FROM t)
         positions = leaves_map[field_name]
-        projection_positions.extend(positions)
         for leaf_position in positions:
             fb_field_index.Start(builder)
             fb_field_index.AddPosition(builder, leaf_position)
@@ -2222,8 +2181,6 @@ def build_query_data_request(schema: 'pa.Schema' = pa.schema([]), predicate: ibi
         builder.PrependUOffsetTRelative(offset)
     projection = builder.EndVector()
 
-    response_schema = VastDBResponseSchema(schema, projection_positions, output_field_names=output_field_names)
-
     fb_source.Start(builder)
     fb_source.AddName(builder, source_name)
     fb_source.AddSchema(builder, schema_obj)
@@ -2237,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':