import logging
import struct
import urllib.parse
from collections import defaultdict, namedtuple
from datetime import datetime
from enum import Enum
from typing import List, Union, Optional, Iterator
import xmltodict
import copy

import flatbuffers
import pyarrow as pa
import requests
import datetime
import hashlib
import hmac
import itertools
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
import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Call as fb_call
import vast_flatbuf.org.apache.arrow.computeir.flatbuf.DateLiteral as fb_date32_lit
import vast_flatbuf.org.apache.arrow.computeir.flatbuf.DecimalLiteral as fb_decimal_lit
import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Expression as fb_expression
import vast_flatbuf.org.apache.arrow.computeir.flatbuf.FieldIndex as fb_field_index
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.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
import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Source as fb_source
import vast_flatbuf.org.apache.arrow.computeir.flatbuf.StringLiteral as fb_string_lit
import vast_flatbuf.org.apache.arrow.computeir.flatbuf.TimeLiteral as fb_time_lit
import vast_flatbuf.org.apache.arrow.computeir.flatbuf.TimestampLiteral as fb_timestamp_lit
import vast_flatbuf.org.apache.arrow.flatbuf.Binary as fb_binary
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.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.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.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.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.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.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

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
ESTORE_INVALID_EHANDLE = UINT64_MAX

"""
S3 Tabular API
"""


def get_logger(name):
    log = logging.getLogger(name)
    log.setLevel(logging.ERROR)
    ch = logging.StreamHandler()
    ch.setLevel(logging.INFO)
    ch.set_name('tabular_stream_handler')
    formatter = logging.Formatter("%(asctime)s:%(levelname)s:%(message)s")
    ch.setFormatter(formatter)
    log.addHandler(ch)
    log.propagate = False
    return log


_logger = get_logger(__name__)


def set_tabular_log_level(level: int = logging.INFO):
    _logger.setLevel(level)


class AuthType(Enum):
    SIGV4 = "s3v4"
    SIGV2 = "s3"
    BASIC = "basic"


class Predicate:
    unit_to_epoch = {
        'ns': 1_000_000,
        'us': 1_000,
        'ms': 1,
        's': 0.001
    }
    unit_to_flatbuff_time_unit = {
        'ns': TimeUnit.NANOSECOND,
        'us': TimeUnit.MICROSECOND,
        'ms': TimeUnit.MILLISECOND,
        's': TimeUnit.SECOND
    }

    def __init__(self, schema: 'pa.Schema', filters: dict):
        self.schema = schema
        self.filters = filters
        self.builder = None
        self._field_name_per_index = None

    def get_field_indexes(self, field: 'pa.Field', field_name_per_index: list) -> None:
        field_name_per_index.append(field.name)

        if isinstance(field.type, pa.StructType):
            flat_fields = field.flatten()
        elif isinstance(field.type, pa.MapType):
            flat_fields = [pa.field(f'{field.name}.entries', pa.struct([field.type.key_field, field.type.item_field]))]
        elif isinstance(field.type, pa.ListType):
            flat_fields = [pa.field(f'{field.name}.{field.type.value_field.name}', field.type.value_field.type)]
        else:
            return

        for flat_field in flat_fields:
            self.get_field_indexes(flat_field, field_name_per_index)

    @property
    def field_name_per_index(self):
        if self._field_name_per_index is None:
            _field_name_per_index = []
            for field in self.schema:
                self.get_field_indexes(field, _field_name_per_index)
            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:
            field_names = self.field_name_per_index.keys()
        projection_fields = []
        for field_name in field_names:
            fb_field_index.Start(builder)
            fb_field_index.AddPosition(builder, self.field_name_per_index[field_name])
            offset = fb_field_index.End(builder)
            projection_fields.append(offset)
        fb_source.StartProjectionVector(builder, len(projection_fields))
        for offset in reversed(projection_fields):
            builder.PrependUOffsetTRelative(offset)
        return builder.EndVector()

    def serialize(self, builder: 'flatbuffers.builder.Builder'):
        self.builder = builder
        offsets = []
        for field_name in self.filters:
            offsets.append(self.build_domain(self.build_column(self.field_name_per_index[field_name]), field_name))
        return self.build_and(offsets)

    def build_column(self, position: int):
        fb_field_index.Start(self.builder)
        fb_field_index.AddPosition(self.builder, position)
        index = fb_field_index.End(self.builder)

        fb_field_ref.Start(self.builder)
        fb_field_ref.AddRefType(self.builder, Deref.FieldIndex)
        fb_field_ref.AddRef(self.builder, index)
        ref = fb_field_ref.End(self.builder)

        fb_expression.Start(self.builder)
        fb_expression.AddImplType(self.builder, ExpressionImpl.FieldRef)
        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]
        _logger.info(f'trying to append field: {field} with domains: {filters}')
        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,
            'ge': self.build_greater_equal,
            'gt': self.build_greater,
            'lt': self.build_less,
            'le': self.build_less_equal,
            'is_not_null': self.build_is_not_null,
            'is_null': self.build_is_null
        }
        if raw_rule in ('is_not_null', 'is_null'):
            # handle is_null or is_valid rule
            op = raw_rule
            value = None
        else:
            op, value = raw_rule.split(maxsplit=1)
        return operator_matcher[op], value

    def build_range(self, column: int, field: pa.Field, filter_by_name: Union[tuple, str]):
        if isinstance(filter_by_name, str):
            filter_by_name = (filter_by_name,)
        if isinstance(filter_by_name, tuple) and len(filter_by_name) == 1:
            op, value = self.rule_to_operator(filter_by_name[0])
            if value:
                literal = self.build_literal(field=field, value=value)
                return op(column, literal)
            return op(column)  # is_null or is_not_null operation

        rules = []
        for rule in filter_by_name:
            op, value = self.rule_to_operator(rule)
            literal = self.build_literal(field=field, value=value)
            rules.append(op(column, literal))

        return self.build_and(rules)

    def build_function(self, name: str, *offsets):
        _logger.info(f'name: {name}, offsets: {offsets}')
        offset_name = self.builder.CreateString(name)
        fb_call.StartArgumentsVector(self.builder, len(offsets))
        for offset in reversed(offsets):
            _logger.info(f'offset: {offset}')
            self.builder.PrependUOffsetTRelative(offset)
        offset_arguments = self.builder.EndVector()

        fb_call.Start(self.builder)
        fb_call.AddName(self.builder, offset_name)
        fb_call.AddArguments(self.builder, offset_arguments)

        offset_call = fb_call.CallEnd(self.builder)

        fb_expression.Start(self.builder)
        fb_expression.AddImplType(self.builder, ExpressionImpl.Call)
        fb_expression.AddImpl(self.builder, offset_call)
        return fb_expression.End(self.builder)

    def build_literal(self, field: pa.Field, value: str):
        if field.type.equals(pa.int64()):
            literal_type = fb_int64_lit
            literal_impl = LiteralImpl.Int64Literal

            field_type_type = Type.Int
            fb_int.Start(self.builder)
            fb_int.AddBitWidth(self.builder, field.type.bit_width)
            fb_int.AddIsSigned(self.builder, True)
            field_type = fb_int.End(self.builder)

            value = int(value)
        elif field.type.equals(pa.int32()):
            literal_type = fb_int32_lit
            literal_impl = LiteralImpl.Int32Literal

            field_type_type = Type.Int
            fb_int.Start(self.builder)
            fb_int.AddBitWidth(self.builder, field.type.bit_width)
            fb_int.AddIsSigned(self.builder, True)
            field_type = fb_int.End(self.builder)

            value = int(value)
        elif field.type.equals(pa.int16()):
            literal_type = fb_int16_lit
            literal_impl = LiteralImpl.Int16Literal

            field_type_type = Type.Int
            fb_int.Start(self.builder)
            fb_int.AddBitWidth(self.builder, field.type.bit_width)
            fb_int.AddIsSigned(self.builder, True)
            field_type = fb_int.End(self.builder)

            value = int(value)
        elif field.type.equals(pa.int8()):
            literal_type = fb_int8_lit
            literal_impl = LiteralImpl.Int8Literal

            field_type_type = Type.Int
            fb_int.Start(self.builder)
            fb_int.AddBitWidth(self.builder, field.type.bit_width)
            fb_int.AddIsSigned(self.builder, True)
            field_type = fb_int.End(self.builder)

            value = int(value)
        elif field.type.equals(pa.float32()):
            literal_type = fb_float32_lit
            literal_impl = LiteralImpl.Float32Literal

            field_type_type = Type.FloatingPoint
            fb_floating_point.Start(self.builder)
            fb_floating_point.AddPrecision(self.builder, 1)  # single
            field_type = fb_floating_point.End(self.builder)

            value = float(value)
        elif field.type.equals(pa.float64()):
            literal_type = fb_float64_lit
            literal_impl = LiteralImpl.Float64Literal

            field_type_type = Type.FloatingPoint
            fb_floating_point.Start(self.builder)
            fb_floating_point.AddPrecision(self.builder, 2)  # double
            field_type = fb_floating_point.End(self.builder)

            value = float(value)
        elif field.type.equals(pa.string()):
            literal_type = fb_string_lit
            literal_impl = LiteralImpl.StringLiteral

            field_type_type = Type.Utf8
            fb_utf8.Start(self.builder)
            field_type = fb_utf8.End(self.builder)

            value = self.builder.CreateString(value)
        elif field.type.equals(pa.date32()):  # pa.date64()
            literal_type = fb_date32_lit
            literal_impl = LiteralImpl.DateLiteral

            field_type_type = Type.Date
            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_value = datetime.strptime(value, '%Y-%m-%d').date()
            date_delta = date_value - start_date
            value = date_delta.days
        elif isinstance(field.type, pa.TimestampType):
            literal_type = fb_timestamp_lit
            literal_impl = LiteralImpl.TimestampLiteral

            field_type_type = Type.Timestamp
            fb_timestamp.Start(self.builder)
            fb_timestamp.AddUnit(self.builder, self.unit_to_flatbuff_time_unit[field.type.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')):

            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]

            field_type_type = Type.Time
            fb_time.Start(self.builder)
            fb_time.AddBitWidth(self.builder, field.type.bit_width)
            fb_time.AddUnit(self.builder, self.unit_to_flatbuff_time_unit[unit])
            field_type = fb_time.End(self.builder)

            value = int(value) * self.unit_to_epoch[unit]
        elif field.type.equals(pa.bool_()):
            literal_type = fb_bool_lit
            literal_impl = LiteralImpl.BooleanLiteral

            field_type_type = Type.Bool
            fb_bool.Start(self.builder)
            field_type = fb_bool.End(self.builder)

            value = True if value == 'true' else False  # not cover all cases
        elif isinstance(field.type, pa.Decimal128Type):
            literal_type = fb_decimal_lit
            literal_impl = LiteralImpl.DecimalLiteral

            field_type_type = Type.Decimal
            fb_decimal.Start(self.builder)
            fb_decimal.AddPrecision(self.builder, field.type.precision)
            fb_decimal.AddScale(self.builder, field.type.scale)
            field_type = fb_decimal.End(self.builder)
            int_value = int(float(value) * 10 ** field.type.scale)
            binary_value = int_value.to_bytes(16, 'little')

            value = self.builder.CreateByteVector(binary_value)
        elif field.type.equals(pa.binary()):
            literal_type = fb_binary_lit
            literal_impl = LiteralImpl.BinaryLiteral

            field_type_type = Type.Binary
            fb_binary.Start(self.builder)
            field_type = fb_binary.End(self.builder)

            value = self.builder.CreateByteVector(value.encode())
        else:
            raise ValueError(f'unsupported predicate for type={field.type}, value={value}')

        literal_type.Start(self.builder)
        literal_type.AddValue(self.builder, value)
        buffer_value = literal_type.End(self.builder)

        fb_field.Start(self.builder)
        fb_field.AddTypeType(self.builder, field_type_type)
        fb_field.AddType(self.builder, field_type)
        buffer_field = fb_field.End(self.builder)

        fb_literal.Start(self.builder)
        fb_literal.AddImplType(self.builder, literal_impl)
        fb_literal.AddImpl(self.builder, buffer_value)
        fb_literal.AddType(self.builder, buffer_field)
        buffer_literal = fb_literal.End(self.builder)

        fb_expression.Start(self.builder)
        fb_expression.AddImplType(self.builder, ExpressionImpl.Literal)
        fb_expression.AddImpl(self.builder, buffer_literal)
        return fb_expression.End(self.builder)

    def build_or(self, offsets: list):
        return self.build_function('or', *offsets)

    def build_and(self, offsets: list):
        return self.build_function('and', *offsets)

    def build_equal(self, column: int, literal: int):
        return self.build_function('equal', column, literal)

    def build_greater(self, column: int, literal: int):
        return self.build_function('greater', column, literal)

    def build_greater_equal(self, column: int, literal: int):
        return self.build_function('greater_equal', column, literal)

    def build_less(self, column: int, literal: int):
        return self.build_function('less', column, literal)

    def build_less_equal(self, column: int, literal: int):
        return self.build_function('less_equal', column, literal)

    def build_is_null(self, column: int):
        return self.build_function('is_null', column)

    def build_is_not_null(self, column: int):
        return self.build_function('is_valid', column)


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.field = field
        self.type = field.type
        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]
        elif isinstance(self.type, pa.ListType):
            self.children = [FieldNode(self.type.value_field, index_iter, parent=self)]
        elif isinstance(self.type, pa.MapType):
            # Map is represented as List<Struct<K, V>> in Arrow
            # TODO: this is a workaround for PyArrow 6 (which has no way to access the internal "entries" Map's field)
            # when we upgrade to latest PyArrow, all nested types can be enumerated using:
            #
            #   if isinstance(self.type, (pa.ListType,  pa.StructType, pa.MapType)):
            #       self.children = [FieldNode(self.type.field(i), index_iter, parent=self) for i in range(self.type.num_fields)]
            #
            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)

    def _iter_to_root(self) -> Iterator['FieldNode']:
        yield self
        if self.parent is not None:
            yield from self.parent._iter_to_root()

    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)
        for child in self.children:
            child.debug_log(level=level+1)

    def set(self, arr: pa.Array):
        """
        Assign the relevant Arrow buffers from the received array into this node.

        VAST can send only a single column at a time - together with its nesting levels.
        For example, `List<List<Int32>>` can be sent in a single go.

        Sending `Struct<a: Int32, b: Float64>` is not supported today, so each column is sent separately
        (possibly duplicating the nesting levels above it).
        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
        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)
        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)}')

    def build(self) -> 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)
        if self.debug:
            _logger.debug('%s result=%s', self.field, result)
        return result


class QueryDataParser:
    """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
        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.leaf_offset = 0

    def parse(self, column: pa.Array):
        """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 < '
                             f'than len(self.leaves): {len(self.leaves)}')
        leaf = self.leaves[self.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()`.
        column = column.combine_chunks()

        if self.debug:
            _logger.debug("parse: index=%d buffers=%s", leaf.index, [b and b.hex() for b in column.buffers()])
        node_list = list(leaf._iter_to_root())  # a FieldNode list (from leaf to root)
        node_list.reverse()  # (from root to leaf)

        # Collect all nesting levels, e.g. `List<Struct<A>` will be split into [`List`, `Struct`, `A`].
        # This way we can extract the buffers from each level and assign them to the appropriate node in loop below.
        array_list = list(_iter_nested_arrays(column))
        if len(array_list) != len(node_list):
            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)

        self.leaf_offset += 1

    def build(self) -> Optional[pa.Table]:
        """Try to build the resulting Table object (if all columns were parsed)"""
        if self.leaf_offset < len(self.leaves):
            return None

        if self.debug:
            for node in self.nodes:
                node.debug_log()

        result = pa.Table.from_arrays(
            arrays=[node.build() 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
    if isinstance(column.type, pa.StructType):
        if not column.type.num_fields == 1:  # Note: VAST serializes only a single struct field at a time
            raise ValueError(f'column.type.num_fields: {column.type.num_fields} not eq to 1')
        yield from _iter_nested_arrays(column.field(0))
    elif isinstance(column.type, pa.ListType):
        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')
def _parse_table_info(obj):

    name = obj.Name().decode()
    properties = obj.Properties().decode()
    handle = obj.Handle().decode()
    num_rows = obj.NumRows()
    used_bytes = obj.SizeInBytes()
    return TableInfo(name, properties, handle, num_rows, used_bytes)

def build_record_batch(column_info, column_values):
    _logger.info(f"column_info={column_info}")
    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()

class VastdbApi:
    def __init__(self, host, access_key, secret_key, username=None, password=None, port=None,
                 secure=False, auth_type=AuthType.SIGV4):
        username = username or ''
        password = password or ''
        if not port:
            port = 443 if secure else 80

        self.session = requests.Session()
        self.session.verify = False
        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 port != 80 and port != 443:
                self.aws_host = f'{host}:{port}'
            else:
                self.aws_host = f'{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')

        proto = "https" if secure else "http"
        self.url = f"{proto}://{self.aws_host}"

    def update_mgmt_session(self, access_key: str, secret_key: str, auth_type=AuthType.SIGV4):
        if auth_type != AuthType.BASIC:
            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')

    def _api_prefix(self, bucket="", schema="", table="", command="", url_params={}):
        prefix_list = [self.url]
        if len(bucket):
            prefix_list.append(bucket)

        if len(schema):
            prefix_list.append(schema)

        if len(table):
            prefix_list.append(table)

        prefix = '/'.join(prefix_list)

        params_list = []
        if command:
            prefix += '?'
            if command != 'list':  # S3 listing request has no command. It looks like 'GET /bucket_name?list-type=2&delimite=Delimiter...'
                params_list.append(command)

        # Query string values must be URL-encoded
        for k, v in sorted(url_params.items()):
            params_list.append(f"{k}={urllib.parse.quote_plus(v)}")

        prefix += '&'.join(params_list)
        return prefix

    def _fill_common_headers(self, txid=0, client_tags=[], version_id=1):
        common_headers = {'tabular-txid': str(txid), 'tabular-api-version-id': str(version_id),
                          'tabular-client-name': 'tabular-api'}
        for tag in client_tags:
            common_headers['tabular-client-tags-%d' % client_tags.index(tag)] = tag

        return common_headers

    def _check_res(self, res, cmd="", expected_retvals=[]):
        try:
            res.raise_for_status()
            if res.status_code != 200:
                if not  res.status_code in expected_retvals:
                    raise ValueError(f"Expected status code mismatch. status_code={res.status_code}")
            else:
                if not len(expected_retvals) == 0:
                    raise ValueError(f"Expected {expected_retvals} but status_code={res.status_code}")
            return res
        except requests.HTTPError as e:
            if res.status_code in expected_retvals:
                _logger.info(f"{cmd} has failed as expected res={res}")
                return res
            else:
                raise e

    def create_schema(self, bucket, name, txid=0, client_tags=[], schema_properties="", expected_retvals=[]):
        """
        Create a collection of tables, use the following request
        POST /bucket/schema?schema HTTP/1.1

        The body of the POST request contains schema properties as flatbuffer
        Request Headers
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag
        """
        builder = flatbuffers.Builder(1024)

        properties_offset = builder.CreateString(schema_properties)
        tabular_create_schema.Start(builder)
        tabular_create_schema.AddProperties(builder, properties_offset)
        params = tabular_create_schema.End(builder)
        builder.Finish(params)
        create_schema_req = builder.Output()

        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        headers['Content-Length'] = str(len(create_schema_req))
        res = self.session.post(self._api_prefix(bucket=bucket, schema=name, command="schema"),
                                data=create_schema_req, headers=headers, stream=True)

        return self._check_res(res, "create_schema", expected_retvals)

    def alter_schema(self, bucket, name, txid=0, client_tags=[], schema_properties="", new_name="", expected_retvals=[]):
        """
        ALTER an existing schema name and/or properties
        PUT /bucket/schema?schema&tabular-new-schema-name=NewSchemaName HTTP/1.1

        Request Headers
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag
        """
        builder = flatbuffers.Builder(1024)

        properties = builder.CreateString(schema_properties)
        tabular_alter_schema.Start(builder)
        if len(schema_properties):
            tabular_alter_schema.AddProperties(builder, properties)

        params = tabular_alter_schema.End(builder)
        builder.Finish(params)
        alter_schema_req = builder.Output()

        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        headers['Content-Length'] = str(len(alter_schema_req))
        url_params = {'tabular-new-schema-name': new_name} if len(new_name) else {}

        res = self.session.put(self._api_prefix(bucket=bucket, schema=name, command="schema", url_params=url_params),
                               data=alter_schema_req, headers=headers)

        return self._check_res(res, "alter_schema", expected_retvals)

    def drop_schema(self, bucket, name, txid=0, client_tags=[], expected_retvals=[]):
        """
        Drop (delete) a schema
        DELETE /bucket/schema?schema HTTP/1.1

        Request Headers
        """
        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)

        res = self.session.delete(self._api_prefix(bucket=bucket, schema=name, command="schema"), headers=headers)

        return self._check_res(res, "drop_schema", expected_retvals)

    def list_schemas(self, bucket, schema="", txid=0, client_tags=[], max_keys=1000, next_key=0, name_prefix="",
                     exact_match=False, expected_retvals=[], count_only=False):
        """
        List all schemas
        GET /bucket/schema_path?schema HTTP/1.1

        Request Headers
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag
        tabular-name-prefix: SchemaNamePrefix
        tabular-max-keys: 1000
        tabular-next-key: NextKey (Schema Name for continuation request)

        The List will return the list in flatbuf format
        """
        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        headers['tabular-max-keys'] = str(max_keys)
        headers['tabular-next-key'] = str(next_key)
        if exact_match:
            headers['tabular-name-exact-match'] = name_prefix
        else:
            headers['tabular-name-prefix'] = name_prefix

        headers['tabular-list-count-only'] = str(count_only)

        schemas = []
        res = self.session.get(self._api_prefix(bucket=bucket, schema=schema, command="schema"), headers=headers, stream=True)
        self._check_res(res, "list_schemas", expected_retvals)
        if res.status_code == 200:
            res_headers = res.headers
            next_key = int(res_headers['tabular-next-key'])
            is_truncated = res_headers['tabular-is-truncated'] == 'true'
            flatbuf = b''.join(res.iter_content(chunk_size=128))
            lists = list_schemas.GetRootAs(flatbuf)
            bucket_name = lists.BucketName().decode()
            if not bucket.startswith(bucket_name):
                raise ValueError(f'bucket: {bucket} did not start from {bucket_name}')
            schemas_length = lists.SchemasLength()
            count = int(res_headers['tabular-list-count']) if 'tabular-list-count' in res_headers else schemas_length
            for i in range(schemas_length):
                schema_obj = lists.Schemas(i)
                name = schema_obj.Name().decode()
                properties = schema_obj.Properties().decode()
                schemas.append([name, properties])

            return bucket_name, schemas, next_key, is_truncated, count

    def list_snapshots(self, bucket, max_keys=1000, next_token=None, expected_retvals=None):
        next_token = next_token or ''
        expected_retvals = expected_retvals or []
        url_params = {'list_type': '2', 'prefix': '.snapshot/', '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]

            return snapshots, is_truncated, marker


    def create_table(self, bucket, schema, name, arrow_schema, txid=0, client_tags=[], expected_retvals=[]):
        """
        Create a table, use the following request
        POST /bucket/schema/table?table HTTP/1.1

        Request Headers
        tabular-txid: <integer> TransactionId
        tabular-client-tag: <string> ClientTag

        The body of the POST request contains table column properties as json
        {
            "format": "string",
            "column_names": {"name1":"type1", "name2":"type2", ...},
            "table_properties": {"key1":"val1", "key2":"val2", ...}
        }
        """
        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)

        serialized_schema = arrow_schema.serialize()
        headers['Content-Length'] = str(len(serialized_schema))

        res = self.session.post(self._api_prefix(bucket=bucket, schema=schema, table=name, command="table"),
                                data=serialized_schema, headers=headers)
        return self._check_res(res, "create_table", expected_retvals)

    def get_table_stats(self, bucket, schema, name, txid=0, client_tags=[], expected_retvals=[]):
        """
        GET /mybucket/myschema/mytable?stats HTTP/1.1
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag

        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:
            res_headers = res.headers
            flatbuf = b''.join(res.iter_content(chunk_size=128))
            stats = get_table_stats.GetRootAs(flatbuf)
            num_rows = stats.NumRows()
            size_in_bytes = stats.SizeInBytes()
            return num_rows, size_in_bytes
        return self._check_res(res, "get_table_stats", expected_retvals)


    def alter_table(self, bucket, schema, name, txid=0, client_tags=[], table_properties="",
                    new_name="", expected_retvals=[]):
        """
        PUT /mybucket/myschema/mytable?table HTTP/1.1
        Content-Length: ContentLength
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag

        Request Body Flatbuffer
        Table properties
        """
        builder = flatbuffers.Builder(1024)

        properties = builder.CreateString(table_properties)
        tabular_alter_table.Start(builder)
        if len(table_properties):
            tabular_alter_table.AddProperties(builder, properties)

        params = tabular_alter_table.End(builder)
        builder.Finish(params)
        alter_table_req = builder.Output()

        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        headers['Content-Length'] = str(len(alter_table_req))
        url_params = {'tabular-new-table-name': new_name} if len(new_name) else {}

        res = self.session.put(self._api_prefix(bucket=bucket, schema=schema, table=name, command="table", url_params=url_params),
                               data=alter_table_req, headers=headers)

        return self._check_res(res, "alter_table", expected_retvals)

    def drop_table(self, bucket, schema, name, txid=0, client_tags=[], expected_retvals=[]):
        """
        DELETE /mybucket/schema_path/mytable?table HTTP/1.1
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag
        """
        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)

        res = self.session.delete(self._api_prefix(bucket=bucket, schema=schema, table=name, command="table"), headers=headers)
        return self._check_res(res, "drop_table", expected_retvals)

    def list_tables(self, bucket, schema, txid=0, client_tags=[], max_keys=1000, next_key=0, name_prefix="",
                    exact_match=False, expected_retvals=[], include_list_stats=False, count_only=False):
        """
        GET /mybucket/schema_path?table HTTP/1.1
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag
        x-tabluar-name-prefix: TableNamePrefix
        tabular-max-keys: 1000
        tabular-next-key: NextKey (Name)
        """
        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        headers['tabular-max-keys'] = str(max_keys)
        headers['tabular-next-key'] = str(next_key)
        if exact_match:
            headers['tabular-name-exact-match'] = name_prefix
        else:
            headers['tabular-name-prefix'] = name_prefix

        headers['tabular-list-count-only'] = str(count_only)
        headers['tabular-include-list-stats'] = str(include_list_stats)

        tables = []
        res = self.session.get(self._api_prefix(bucket=bucket, schema=schema, command="table"), headers=headers)
        self._check_res(res, "list_table", expected_retvals)
        if res.status_code == 200:
            res_headers = res.headers
            next_key = int(res_headers['tabular-next-key'])
            is_truncated = res_headers['tabular-is-truncated'] == 'true'
            flatbuf = b''.join(res.iter_content(chunk_size=128))
            lists = list_tables.GetRootAs(flatbuf)
            bucket_name = lists.BucketName().decode()
            schema_name = lists.SchemaName().decode()
            if not bucket.startswith(bucket_name):  # ignore snapshot name
                raise ValueError(f'bucket: {bucket} did not start from {bucket_name}')
            tables_length = lists.TablesLength()
            count = int(res_headers['tabular-list-count']) if 'tabular-list-count' in res_headers else tables_length
            for i in range(tables_length):
                tables.append(_parse_table_info(lists.Tables(i)))

            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
        POST /bucket/schema/table?column HTTP/1.1

        Request Headers
        tabular-txid: <integer> TransactionId
        tabular-client-tag: <string> ClientTag

        The body of the POST request contains table column properties as json
        {
            "format": "string",
            "column_names": {"name1":"type1", "name2":"type2", ...},
            "table_properties": {"key1":"val1", "key2":"val2", ...}
        }
        """
        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        serialized_schema = arrow_schema.serialize()
        headers['Content-Length'] = str(len(serialized_schema))

        res = self.session.post(self._api_prefix(bucket=bucket, schema=schema, table=name, command="column"),
                                data=serialized_schema, headers=headers)
        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=[]):
        """
        PUT /bucket/schema/table?column&tabular-column-name=ColumnName&tabular-new-column-name=NewColumnName HTTP/1.1
        Content-Length: ContentLength
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag
        tabular-column-sep: ColumnSep

        Request Body Flatbuffer
        Table properties
        """
        builder = flatbuffers.Builder(1024)

        properties = builder.CreateString(column_properties)
        stats = builder.CreateString(column_stats)
        tabular_alter_column.Start(builder)
        if len(column_properties):
            tabular_alter_column.AddProperties(builder, properties)
        if len(column_stats):
            tabular_alter_column.AddStats(builder, stats)

        params = tabular_alter_column.End(builder)
        builder.Finish(params)
        alter_column_req = builder.Output()

        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        headers['tabular-column-sep'] = column_sep
        headers['Content-Length'] = str(len(alter_column_req))

        url_params = {'tabular-column-name': name }
        if len(new_name):
            url_params['tabular-new-column-name'] = new_name

        res = self.session.put(self._api_prefix(bucket=bucket, schema=schema, table=table, command="column", url_params=url_params),
                               data=alter_column_req, headers=headers)
        return self._check_res(res, "alter_column", expected_retvals)

    def drop_columns(self, bucket, schema, table, arrow_schema, txid=0, client_tags=[], expected_retvals=[]):
        """
        DELETE /mybucket/myschema/mytable?column HTTP/1.1
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag
        tabular-column-name: OldColumnName
        """
        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        serialized_schema = arrow_schema.serialize()
        headers['Content-Length'] = str(len(serialized_schema))

        res = self.session.delete(self._api_prefix(bucket=bucket, schema=schema, table=table, command="column"),
                                data=serialized_schema, headers=headers)
        return self._check_res(res, "drop_columns", expected_retvals)

    def list_columns(self, bucket, schema, table, *, txid=0, client_tags=None, max_keys=1000, next_key=0,
                     count_only=False, name_prefix="", exact_match=False,
                     expected_retvals=None, bc_list_internals=False):
        """
        GET /mybucket/myschema/mytable?columns HTTP/1.1
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag
        x-tabluar-name-prefix: TableNamePrefix
        tabular-max-keys: 1000
        tabular-next-key: NextColumnId
        """
        client_tags = client_tags or []
        expected_retvals = expected_retvals or []

        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        headers['tabular-max-keys'] = str(max_keys)
        headers['tabular-next-key'] = str(next_key)
        headers['tabular-list-count-only'] = str(count_only)
        if bc_list_internals:
            headers['tabular-bc-list-internal-col'] = "true"

        res = self.session.get(self._api_prefix(bucket=bucket, schema=schema, table=table, command="column"),
                               headers=headers, stream=True)
        self._check_res(res, "list_columns", expected_retvals)
        if res.status_code == 200:
            res_headers = res.headers
            next_key = int(res_headers['tabular-next-key'])
            is_truncated = res_headers['tabular-is-truncated'] == 'true'
            count = int(res_headers['tabular-list-count'])
            columns = []
            if not count_only:
                schema_buf = b''.join(res.iter_content(chunk_size=128))
                schema_out = pa.ipc.open_stream(schema_buf).schema
    #            _logger.info(f"schema={schema_out}")
                for i in range(len(schema_out)):
                    f = schema_out.field(i)
                    add_column = not name_prefix or (exact_match and f.name == name_prefix) or (not exact_match and f.name.startswith(name_prefix))
                    if add_column:
                        columns.append([f.name, f.type, f.metadata])

            return columns, next_key, is_truncated, count

    def begin_transaction(self, client_tags=[], expected_retvals=[]):
        """
        POST /?transaction HTTP/1.1
        tabular-client-tag: ClientTag

        Response
        tabular-txid: TransactionId
        """
        headers = self._fill_common_headers(client_tags=client_tags)
        res = self.session.post(self._api_prefix(command="transaction"), headers=headers)
        return self._check_res(res, "begin_transaction", expected_retvals)

    def commit_transaction(self, txid, client_tags=[], expected_retvals=[]):
        """
        PUT /?transaction HTTP/1.1
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag
        """
        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        res = self.session.put(self._api_prefix(command="transaction"), headers=headers)
        return self._check_res(res, "commit_transaction", expected_retvals)

    def rollback_transaction(self, txid, client_tags=[], expected_retvals=[]):
        """
        DELETE /?transaction HTTP/1.1
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag
        """
        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        res = self.session.delete(self._api_prefix(command="transaction"), headers=headers)
        return self._check_res(res, "rollback_transaction", expected_retvals)

    def get_transaction(self, txid, client_tags=[], expected_retvals=[]):
        """
        GET /?transaction HTTP/1.1
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag
        """
        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        res = self.session.get(self._api_prefix(command="transaction"), headers=headers)
        return self._check_res(res, "get_transaction", expected_retvals)

    def query_data(self, bucket, schema, table, params, split=(0, 1, 8), num_sub_splits=1, response_row_id=False,
                   txid=0, client_tags=[], expected_retvals=[], limit_rows=0, schedule_id=None, retry_count=0,
                   search_path=None, sub_split_start_row_ids=[], tenant_guid=None, projection='', enable_sorted_projections=True):
        """
        GET /mybucket/myschema/mytable?data HTTP/1.1
        Content-Length: ContentLength
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag
        tabular-split: "split_id,total_splits,num_row_groups_per_split"
        tabular-num-of-subsplits: "total"
        tabular-request-format: "string"
        tabular-response-format: "string" #arrow/trino
        tabular-response-row-id: "true" # default false
        tabular-schedule-id: "schedule-id"

        Request Body (flatbuf)
        projections_chunk [expressions]
        predicate_chunk "formatted_data", (required)

        """
        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        headers['Content-Length'] = str(len(params))
        headers['tabular-split'] = ','.join(map(str, split))
        headers['tabular-num-of-subsplits'] = str(num_sub_splits)
        headers['tabular-request-format'] = "string"
        headers['tabular-response-format'] = "string"
        headers['tabular-enable-sorted-projections'] = str(enable_sorted_projections)
        if response_row_id:
            headers['tabular-response-row-id'] = "true"
        if limit_rows > 0:
            headers['tabular-limit-rows'] = str(limit_rows)
        if schedule_id is not None:
            headers['tabular-schedule-id'] = str(schedule_id)
        if retry_count > 0:
            headers['tabular-retry-count'] = str(retry_count)
        if search_path:
            headers['tabular-bc-search-path'] = str(search_path)
        if tenant_guid:
            headers['tabular-tenant-guid'] = str(tenant_guid)

        if len(sub_split_start_row_ids) == 0:
            sub_split_start_row_ids = [(n, 0) for n in range(num_sub_splits)]

        for sub_split_id, start_row_id in sub_split_start_row_ids:
            headers[f'tabular-start-row-id-{sub_split_id}'] = f"{sub_split_id},{start_row_id}"

        url_params = {'name': projection} if len(projection) else {}

        res = self.session.get(self._api_prefix(bucket=bucket, schema=schema, table=table, command="data", url_params=url_params),
                               data=params, headers=headers, stream=True)
        return self._check_res(res, "query_data", expected_retvals)


    def _list_table_columns(self, bucket, schema, table, filters=None, field_names=None):
        # build a list of the queried column names
        queried_columns = []
        # get all columns from the table
        all_columns = []
        next_key = 0
        while True:
            cur_columns, next_key, is_truncated, count = self.list_columns(bucket=bucket, schema=schema,
                                                                           table=table,
                                                                           next_key=next_key)
            if not cur_columns:
                break;
            all_columns.extend(cur_columns)
            if not is_truncated:
                break;

        # build a list of the queried columns
        queried_column_names = set()
        if filters:
            queried_column_names.update(filters.keys())
        if field_names:
            queried_column_names.update(field_names)
        all_column_names = set()
        for column in all_columns:
            column_name = column[0]
            all_column_names.add(column_name)
            if column_name in queried_column_names or not field_names:
                queried_columns.append(column)

        # check that all queried column names are in the table
        _logger.info(f"_list_table_columns: all_column_names={all_column_names} filters={filters} field_names={field_names}")
        if filters:
            for filter_column_name in filters.keys():
                if filter_column_name not in all_column_names:
                    raise KeyError((f'filter column name: {filter_column_name} does not appear in the table'))
        if field_names:
            for field_name in field_names:
                if field_name not in all_column_names:
                    raise ValueError((f'field name: {field_name} does not appear in the table'))
        return queried_columns

    def _begin_tx_if_necessary(self, txid):
        if not txid:
            created_txid = True
            res = self.begin_transaction()
            txid = res.headers.get('tabular-txid')
        else:
            created_txid = False

        return txid, created_txid

    def _prepare_query(self, bucket, schema, table, num_sub_splits, filters, field_names):
            queried_columns = self._list_table_columns(bucket, schema, table, filters, field_names)
            arrow_schema = pa.schema([(column[0], column[1]) for column in queried_columns])
            query_data_request = build_query_data_request(schema=arrow_schema, filters=filters, field_names=field_names)
            start_row_ids = {i:0 for i in range(num_sub_splits)}
            return queried_columns, arrow_schema, query_data_request, start_row_ids

    def _more_pages_exist(self, start_row_ids):
        for row_id in start_row_ids.values():
            if row_id != TABULAR_INVALID_ROW_ID:
                return True
        return False

    def _query_page(self, bucket, schema, table, query_data_request, split=(0, 1, 8), num_sub_splits=1, response_row_id=False,
                   txid=0, limit_rows=0, sub_split_start_row_ids=[], filters=None, field_names=None):
        res = self.query_data(bucket=bucket, schema=schema, table=table, params=query_data_request.serialized, split=split,
                              num_sub_splits=num_sub_splits, response_row_id=response_row_id, txid=txid,
                              limit_rows=limit_rows, sub_split_start_row_ids=sub_split_start_row_ids)
        start_row_ids = {}
        subsplit_tables = parse_query_data_response(res.raw, query_data_request.response_schema,
                                                    start_row_ids=start_row_ids)
        table_page = pa.concat_tables(subsplit_tables)
        _logger.info("query_page: table_page num_rows=%s start_row_ids len=%s",
                     len(table_page), len(start_row_ids))

        return table_page, start_row_ids

    def _query_page_iterator(self, bucket, schema, table, query_data_request, split=(0, 1, 8), num_sub_splits=1, response_row_id=False,
                             txid=0, limit_rows=0, start_row_ids={}, filters=None, field_names=None):
        res = self.query_data(bucket=bucket, schema=schema, table=table, params=query_data_request.serialized, split=split,
                              num_sub_splits=num_sub_splits, response_row_id=response_row_id, txid=txid,
                              limit_rows=limit_rows, sub_split_start_row_ids=start_row_ids.items())
        for subsplit_table in parse_query_data_response(res.raw, query_data_request.response_schema,
                                                        start_row_ids=start_row_ids):
            for record_batch in subsplit_table.to_batches():
                yield record_batch
        _logger.info(f"query_page_iterator: start_row_ids={start_row_ids}")

    def query_iterator(self, bucket, schema, table, split=(0, 1, 8), num_sub_splits=1, response_row_id=False,
                       txid=0, limit_per_subsplit=128*1024, filters=None, field_names=None):
        """
        query rows into a table.

        Parameters
        ----------
        bucket : string
            The bucket of the table.
        schema : string
            The schema of the table.
        table : string
            The table name.
        split : tuple
            The tuple consists of 3 integers: (split id, total splits, row groups per subsplit).
            Splits allow parallelization of query on multiple sessions, for example a session per
            cnode
            default: (0,1,8)
        num_sub_splits : integer
            The number of subsplits to use per split. Sub-splits define the level of parallelism per
            single execution of a split inside a cnode.
            default: 1
        response_row_id : boolean
            Return a column with the internal row ids of the table
            default: False
        txid : integer
            A transaction id. The transaction may be initiated before the query, and if not, the query will initiate it
        limit_per_subsplit : integer
            Limit the number of rows per subsplit
            default:131072
        filters : dict
            A dictionary whose keys are column names, and values are lists of string expressions that represent
            filter conditions on the column. AND is applied on the conditions. The condition formats are:
            'column_name eq some_value'
        filters : dict
            A list of column names to be returned in the output table

        Returns
        -------
        Query iterator generator

        Yields
        ------
        pyarrow.RecordBatch

        Examples
        --------
        for record_batch in query_iterator('some_bucket', 'some_schema', 'some_table',
                                           filters={'name': ['eq Alice', 'eq Bob']}
                                           field_names=['name','age']):
            ...

        """


        #create a transaction if necessary
        txid, created_txid = self._begin_tx_if_necessary(txid)

        try:
            # prepare query
            queried_columns, arrow_schema, query_data_request, start_row_ids = \
                self._prepare_query(bucket, schema, table, num_sub_splits, filters, field_names)

            while self._more_pages_exist(start_row_ids):
                for record_batch in self._query_page_iterator(bucket=bucket, schema=schema, table=table, query_data_request=query_data_request,
                                                             split=split, num_sub_splits=num_sub_splits, response_row_id=response_row_id,
                                                             txid=txid, limit_rows=limit_per_subsplit,
                                                             start_row_ids=start_row_ids):
                    yield record_batch

            # commit if needed
            if created_txid:
                self.commit_transaction(txid)

        except Exception as e:
            # rollback if needed
            if created_txid:
                self.rollback_transaction(txid)
            raise e

    def query(self, bucket, schema, table, split=(0, 1, 8), num_sub_splits=1, response_row_id=False,
              txid=0, limit=0, limit_per_subsplit=131072, filters=None, field_names=None, paginated_response=False):
        """
        query rows into a table.

        Parameters
        ----------
        bucket : string
            The bucket of the table.
        schema : string
            The schema of the table.
        table : string
            The table name.
        split : tuple
            The tuple consists of 3 integers: (split id, total splits, row groups per subsplit).
            Splits allow parallelization of query on multiple sessions, for example a session per
            cnode
            default: (0,1,8)
        num_sub_splits : integer
            The number of subsplits to use per split. Sub-splits define the level of parallelism per
            single execution of a split inside a cnode.
            default: 1
        response_row_id : boolean
            Return a column with the internal row ids of the table
            default: False
        txid : integer
            A transaction id. The transaction may be initiated before the query, and be used to provide
            multiple ACID operations
        limit : integer
            Limit the number of rows in the response
            default: 0 (no limit)
        limit_per_subsplit : integer
            Limit the number of rows per subsplit
            default:131072
        filters : dict
            A dictionary whose keys are column names, and values are lists of string expressions that represent
            filter conditions on the column. AND is applied on the conditions. The condition formats are:
            'column_name eq some_value'
        filters : dict
            A list of column names to be returned in the output table

        Returns
        -------
        pyarrow.Table


        Examples
        --------
        table = query('some_bucket', 'some_schema', 'some_table',
                      filters={'name': ['eq Alice', 'eq Bob']}
                      field_names=['name','age'])

        """


        #create a transaction
        txid, created_txid = self._begin_tx_if_necessary(txid)

        try:
            # prepare query
            queried_columns, arrow_schema, query_data_request, start_row_ids = \
                self._prepare_query(bucket, schema, table, num_sub_splits, filters, field_names)

            table_pages = []
            row_count = 0
            while (self._more_pages_exist(start_row_ids) and
                   (not limit or row_count < limit)):
                limit_rows = min(limit_per_subsplit, limit-row_count)
                table_page, start_row_ids = self._query_page(bucket=bucket, schema=schema, table=table, query_data_request=query_data_request,
                                                             split=split, num_sub_splits=num_sub_splits, response_row_id=response_row_id,
                                                             txid=txid, limit_rows=limit_rows,
                                                             sub_split_start_row_ids=start_row_ids.items())
                table_pages.append(table_page)
                row_count += len(table_page)
                _logger.info("query: table_pages len=%s row_count=%s",
                             len(table_pages), row_count)

            # commit if needed
            if created_txid:
                self.commit_transaction(txid)

            # concatenate all table pages and return result
            out_table = pa.concat_tables(table_pages)
            out_table = out_table.slice(length=limit) if limit else out_table
            _logger.info("query: out_table len=%s row_count=%s",
                          len(out_table), len(out_table))

            return out_table
        except Exception as e:
            # rollback if needed
            if created_txid:
                self.rollback_transaction(txid)
            raise e


    """
    source_files: list of (bucket_name, file_name)
    """
    def import_data(self, bucket, schema, table, source_files, txid=0, client_tags=[], expected_retvals=[], case_sensitive=True,
                    schedule_id=None, retry_count=0):
        """
        POST /mybucket/myschema/mytable?data HTTP/1.1
        Content-Length: ContentLength
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag
        tabular-case-sensitive: True (default)

        Request Body flatbuffer
        table S3File {
          format: "parquet";
          bucket_name: "bucket1";
          file_name: "path/to/file1";
          partitions: [ubyte]; # serialized as a Arrow schema + single-row Arrow page, if exist
        }
        table ImportDataRequest {
          s3_files: [S3File]
        }

        """
        builder = flatbuffers.Builder(1024)
        s3_files = []
        for (src_bucket_name, src_obj_name), partition_record_bytes in source_files.items():
            fmt = builder.CreateString('parquet')
            bname = builder.CreateString(src_bucket_name)
            fname = builder.CreateString(src_obj_name)
            partition = builder.CreateByteVector(partition_record_bytes)
            tabular_s3_file.Start(builder)
            tabular_s3_file.AddFormat(builder, fmt)
            tabular_s3_file.AddBucketName(builder, bname)
            tabular_s3_file.AddFileName(builder, fname)
            if len(partition_record_bytes) > 0:
                tabular_s3_file.AddPartitions(builder, partition)

            s3_file = tabular_s3_file.End(builder)
            s3_files.append(s3_file)

        tabular_import_data.StartS3FilesVector(builder, len(s3_files))
        for f in reversed(s3_files):
            builder.PrependUOffsetTRelative(f)

        files = builder.EndVector()
        tabular_import_data.Start(builder)
        tabular_import_data.AddS3Files(builder, files)
        params = tabular_import_data.End(builder)
        builder.Finish(params)
        import_req = builder.Output()

        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        headers['Content-Length'] = str(len(import_req))
        headers['tabular-case-sensitive'] = str(case_sensitive)
        if schedule_id is not None:
            headers['tabular-schedule-id'] = str(schedule_id)
        if retry_count > 0:
            headers['tabular-retry-count'] = str(retry_count)
        res = self.session.post(self._api_prefix(bucket=bucket, schema=schema, table=table, command="data"),
                                data=import_req, headers=headers, stream=True)

        return self._check_res(res, "import_data", expected_retvals)

    def merge_data(self):
        """
        TODO

        POST /mybucket/myschema/mytable?data HTTP/1.1
        Content-Length: ContentLength
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag

        Request Body
        {
          "format": "string",
          "select_source": "formatted data"
          "predicate": "formatted_data"
        }
        """
        pass

    def insert(self, bucket, schema, table, rows):
        """
        insert rows into a table.

        Parameters
        ----------
        bucket : string
            The bucket of the table.
        schema : string
            The schema of the table.
        table : string
            The table name.

        rows : dict
            The rows to insert.
            dictionary key: column name
            dictionary value: array of cell values to insert

        Returns
        -------
        None


        Examples
        --------
        insert('some_bucket', 'some_schema', 'some_table', {'name': ['Alice','Bob'], 'age': [25,24]})

        """
        columns = self._list_table_columns(bucket, schema, table, field_names=rows.keys())
        columns_dict = dict([(column[0], column[1]) for column in columns])
        arrow_schema = pa.schema([])
        arrays = []
        for column_name, column_values in rows.items():
            column_type = columns_dict[column_name]
            field = pa.field(column_name, column_type)
            arrow_schema = arrow_schema.append(field)
            arrays.append(pa.array(column_values, column_type))
        batch = pa.record_batch(arrays, arrow_schema)
        insert_rows_req = serialize_record_batch(batch)
        self.insert_rows(bucket=bucket, schema=schema,
                         table=table, record_batch=insert_rows_req)

    def insert_rows(self, bucket, schema, table, record_batch, txid=0, client_tags=[], expected_retvals=[]):
        """
        POST /mybucket/myschema/mytable?rows HTTP/1.1
        Content-Length: ContentLength
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag

        Request Body
            RecordBatch
        """
        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        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)

    def update_rows(self, bucket, schema, table, record_batch, txid=0, client_tags=[], expected_retvals=[]):
        """
        PUT /mybucket/myschema/mytable?rows HTTP/1.1
        Content-Length: ContentLength
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag

        Request Body
            RecordBatch where first column must be row_id
        """
        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        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)

    def delete_rows(self, bucket, schema, table, record_batch, txid=0, client_tags=[], expected_retvals=[]):
        """
        DELETE /mybucket/myschema/mytable?rows HTTP/1.1
        Content-Length: ContentLength
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag

        Request Body
            RecordBatch with single column '$row_id'
        """
        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)

    def create_projection(self, bucket, schema, table, name, columns, txid=0, client_tags=[], expected_retvals=[]):
        """
        Create a table, use the following request
        POST /bucket/schema/table?projection&name=my_projection HTTP/1.1

        The user may define several Sorted and/Or Unsorted columns from the original table
        columns = [('col1','Sorted'), ('col2','Unsorted'), ('col3','Unsorted')]

        Request Headers
        tabular-txid: <integer> TransactionId
        tabular-client-tag: <string> ClientTag

        The body of the POST request contains table column properties as Flatbuffer
        {
            column_type: Sorted, Unsorted
        }
        """
        builder = flatbuffers.Builder(1024)
        proj_columns = []
        for (col_name, col_type) in columns:
            pname = builder.CreateString(col_name)
            ptype = tabular_proj_column_type.ColumnType.Sorted if col_type == "Sorted" else tabular_proj_column_type.ColumnType.Unsorted
            tabular_projecion_column.Start(builder)
            tabular_projecion_column.AddName(builder, pname)
            tabular_projecion_column.AddType(builder, ptype)

            column_info = tabular_projecion_column.End(builder)
            proj_columns.append(column_info)

        tabular_create_projection.StartColumnsVector(builder, len(proj_columns))
        for c in reversed(proj_columns):
            builder.PrependUOffsetTRelative(c)

        columns = builder.EndVector()
        tabular_create_projection.Start(builder)
        tabular_create_projection.AddColumns(builder, columns)
        params = tabular_create_projection.End(builder)
        builder.Finish(params)
        create_projection_req = builder.Output()

        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)

        headers['Content-Length'] = str(len(create_projection_req))
        url_params = {'name': name}

        res = self.session.post(self._api_prefix(bucket=bucket, schema=schema, table=table, command="projection", url_params=url_params),
                                data=create_projection_req, headers=headers)
        return self._check_res(res, "create_projection", expected_retvals)

    def get_projection_stats(self, bucket, schema, table, name, txid=0, client_tags=[], expected_retvals=[]):
        """
        GET /mybucket/myschema/mytable?projection-stats&name=my_projection HTTP/1.1
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag

        The Command will return the statistics in flatbuf format
        """
        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        url_params = {'name': name}
        res = self.session.get(self._api_prefix(bucket=bucket, schema=schema, table=table, command="projection-stats", url_params=url_params),
                               headers=headers)
        if res.status_code == 200:
            flatbuf = b''.join(res.iter_content(chunk_size=128))
            stats = get_projection_table_stats.GetRootAs(flatbuf)
            num_rows = stats.NumRows()
            size_in_bytes = stats.SizeInBytes()
            dirty_blocks_percentage = stats.DirtyBlocksPercentage()
            initial_sync_progress = stats.InitialSyncProgress()
            return num_rows, size_in_bytes, dirty_blocks_percentage, initial_sync_progress

        return self._check_res(res, "get_projection_stats", expected_retvals)

    def alter_projection(self, bucket, schema, table, name, txid=0, client_tags=[], table_properties="",
                         new_name="", expected_retvals=[]):
        """
        PUT /mybucket/myschema/mytable?projection&name=my_projection HTTP/1.1
        Content-Length: ContentLength
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag

        Request Body Flatbuffer
        Table properties
        """
        builder = flatbuffers.Builder(1024)

        properties = builder.CreateString(table_properties)
        table_name = builder.CreateString(new_name)
        tabular_alter_projection.Start(builder)
        if new_name:
            tabular_alter_projection.AddNewName(builder, table_name)

        if table_properties:
            tabular_alter_projection.AddProperties(builder, properties)

        params = tabular_alter_projection.End(builder)
        builder.Finish(params)
        alter_projection_req = builder.Output()

        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        headers['Content-Length'] = str(len(alter_projection_req))
        url_params = {'name': name}

        res = self.session.put(self._api_prefix(bucket=bucket, schema=schema, table=table, command="projection", url_params=url_params),
                               data=alter_projection_req, headers=headers)

        return self._check_res(res, "alter_projection", expected_retvals)

    def drop_projection(self, bucket, schema, table, name, txid=0, client_tags=[], expected_retvals=[]):
        """
        DELETE /mybucket/schema_path/mytable?projection&name=my_projection HTTP/1.1
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag
        """
        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        url_params = {'name': name}

        res = self.session.delete(self._api_prefix(bucket=bucket, schema=schema, table=table, command="projection", url_params=url_params),
                                  headers=headers)
        return self._check_res(res, "drop_projection", expected_retvals)

    def list_projections(self, bucket, schema, table, txid=0, client_tags=[], max_keys=1000, next_key=0, name_prefix="",
                         exact_match=False, expected_retvals=[], include_list_stats=False, count_only=False):
        """
        GET /mybucket/schema_path/my_table?projection HTTP/1.1
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag
        x-tabluar-name-prefix: TableNamePrefix
        tabular-max-keys: 1000
        tabular-next-key: NextKey (Name)
        """
        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        headers['tabular-max-keys'] = str(max_keys)
        headers['tabular-next-key'] = str(next_key)
        if exact_match:
            headers['tabular-name-exact-match'] = name_prefix
        else:
            headers['tabular-name-prefix'] = name_prefix

        headers['tabular-list-count-only'] = str(count_only)
        headers['tabular-include-list-stats'] = str(include_list_stats)

        projections = []
        res = self.session.get(self._api_prefix(bucket=bucket, schema=schema, table=table, command="projection"), headers=headers)
        self._check_res(res, "list_projections", expected_retvals)
        if res.status_code == 200:
            res_headers = res.headers
            next_key = int(res_headers['tabular-next-key'])
            is_truncated = res_headers['tabular-is-truncated'] == 'true'
            count = int(res_headers['tabular-list-count'])
            flatbuf = b''.join(res.iter_content(chunk_size=128))
            lists = list_projections.GetRootAs(flatbuf)
            bucket_name = lists.BucketName().decode()
            schema_name = lists.SchemaName().decode()
            table_name = lists.TableName().decode()
            if not bucket.startswith(bucket_name):  # ignore snapshot name
                raise ValueError(f'bucket: {bucket} did not start from {bucket_name}')
            projections_length = lists.ProjectionsLength()
            for i in range(projections_length):
                projections.append(_parse_table_info(lists.Projections(i)))

            return bucket_name, schema_name, table_name, projections, next_key, is_truncated, count

    def list_projection_columns(self, bucket, schema, table, projection, txid=0, client_tags=[], max_keys=1000,
                                next_key=0, count_only=False, name_prefix="", exact_match=False,
                                expected_retvals=None):
        """
        GET /mybucket/myschema/mytable?projection-columns HTTP/1.1
        tabular-txid: TransactionId
        tabular-client-tag: ClientTag
        x-tabluar-name-prefix: TableNamePrefix
        tabular-max-keys: 1000
        tabular-next-key: NextColumnId
        """
        client_tags = client_tags or []
        expected_retvals = expected_retvals or []

        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
        headers['tabular-max-keys'] = str(max_keys)
        headers['tabular-next-key'] = str(next_key)
        headers['tabular-list-count-only'] = str(count_only)
        url_params = {'name': projection}

        res = self.session.get(self._api_prefix(bucket=bucket, schema=schema, table=table, command="projection-columns", url_params=url_params),
                               headers=headers, stream=True)
        self._check_res(res, "list_projection_columns", expected_retvals)
        # list projection columns response will also show column type Sorted/UnSorted
        if res.status_code == 200:
            res_headers = res.headers
            next_key = int(res_headers['tabular-next-key'])
            is_truncated = res_headers['tabular-is-truncated'] == 'true'
            count = int(res_headers['tabular-list-count'])
            columns = []
            if not count_only:
                schema_buf = b''.join(res.iter_content(chunk_size=128))
                schema_out = pa.ipc.open_stream(schema_buf).schema
                for i in range(len(schema_out)):
                    f = schema_out.field(i)
                #   sort_type = f.metadata[b'VAST:sort_type'].decode()
                    add_column = not name_prefix or (exact_match and f.name == name_prefix) or (not exact_match and f.name.startswith(name_prefix))
                    if add_column:
                        columns.append([f.name, f.type, f.metadata])

            return columns, next_key, is_truncated, count


def _iter_query_data_response_columns(fileobj, stream_ids=None):
    readers = {}  # {stream_id: pa.ipc.RecordBatchStreamReader}
    while True:
        stream_id_bytes = fileobj.read(4)
        if not stream_id_bytes:
            if readers:
                raise EOFError(f'no readers ({readers}) should be open at EOF')
            break

        stream_id, = struct.unpack('<L', stream_id_bytes)
        if stream_ids is not None:
            stream_ids.update([stream_id])  # count stream IDs using a collections.Counter
        if stream_id == TABULAR_KEEP_ALIVE_STREAM_ID:
#            _logger.info(f"stream_id={stream_id} (skipping)")
            continue

        if stream_id == TABULAR_QUERY_DATA_COMPLETED_STREAM_ID:
            # read the terminating end chunk from socket
            res = fileobj.read()
            _logger.info(f"stream_id={stream_id} res={res} (finish)")
            return

        if stream_id == TABULAR_QUERY_DATA_FAILED_STREAM_ID:
            # read the terminating end chunk from socket
            res = fileobj.read()
            _logger.info(f"stream_id={stream_id} res={res} (failed)")
            raise IOError(f"Query data stream failed res={res}")

        next_row_id_bytes = fileobj.read(8)
        next_row_id, = struct.unpack('<Q', next_row_id_bytes)
        _logger.info(f"stream_id={stream_id} next_row_id={next_row_id}")

        if stream_id not in readers:
            # we implicitly read 1st message (Arrow schema) when constructing RecordBatchStreamReader
            reader = pa.ipc.RecordBatchStreamReader(fileobj)
            _logger.info(f"stream_id={stream_id} schema={reader.schema}")
            readers[stream_id] = (reader, [])
            continue

        (reader, batches) = readers[stream_id]
        try:
            batch = reader.read_next_batch() # read single-column chunk data
            _logger.info(f"stream_id={stream_id} rows={len(batch)} chunk={batch}")
            batches.append(batch)
        except StopIteration:  # we got an end-of-stream IPC message for a given stream ID
            reader, batches = readers.pop(stream_id)  # end of column
            table = pa.Table.from_batches(batches)  # concatenate all column chunks (as a single)
            _logger.info(f"stream_id={stream_id} rows={len(table)} column={table}")
            yield (stream_id, next_row_id, table)


def parse_query_data_response(conn, schema, stream_ids=None, start_row_ids=None, debug=False):
    """
    Generates pyarrow.Table objects from QueryData API response stream.

    A pyarrow.Table is a helper class that combines a Schema with multiple RecordBatches and allows easy data access.
    """
    if start_row_ids is None:
        start_row_ids = {}
    is_empty_projection = (len(schema) == 0)
    parsers = defaultdict(lambda: QueryDataParser(schema, debug=debug))  # {stream_id: QueryDataParser}
    for stream_id, next_row_id, table in _iter_query_data_response_columns(conn, stream_ids):
        parser = parsers[stream_id]
        for column in table.columns:
            parser.parse(column)

        parsed_table = parser.build()
        if parsed_table is not None:  # when we got all columns (and before starting a new "select_rows" cycle)
            parsers.pop(stream_id)
            if is_empty_projection:  # VAST returns an empty RecordBatch, with the correct rows' count
                parsed_table = table

            _logger.info(f"stream_id={stream_id} rows={len(parsed_table)} next_row_id={next_row_id} table={parsed_table}")
            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}')

def build_field(builder: flatbuffers.Builder, f: pa.Field, name: str):
    _logger.info(f"name={f.name}")
    children = None
    if isinstance(f.type, pa.StructType):
        children = [build_field(builder, child, child.name) for child in list(f.type)]
    if isinstance(f.type, pa.ListType):
        children = [build_field(builder, f.type.value_field, "item")]
    if isinstance(f.type, pa.MapType):
        children = [
            build_field(builder, f.type.key_field, "key"),
            build_field(builder, f.type.item_field, "value"),
        ]

        # adding "entries" column:
        fb_field.StartChildrenVector(builder, len(children))
        for offset in reversed(children):
            builder.PrependUOffsetTRelative(offset)
        children = builder.EndVector()

        child_col_name = builder.CreateString("entries")
        fb_field.Start(builder)
        fb_field.AddName(builder, child_col_name)
        fb_field.AddChildren(builder, children)

        _logger.info(f"added key and map to entries")
        children = [fb_field.End(builder)]

    if children is not None:
        fb_field.StartChildrenVector(builder, len(children))
        for offset in reversed(children):
            builder.PrependUOffsetTRelative(offset)
        children = builder.EndVector()

    col_name = builder.CreateString(name)
    _logger.info(f"add col_name={name} to fb")
    fb_field.Start(builder)
    fb_field.AddName(builder, col_name)
    if children is not None:
        _logger.info(f"add col_name={name} childern")
        fb_field.AddChildren(builder, children)
    return fb_field.End(builder)


class QueryDataRequest:
    def __init__(self, serialized, response_schema):
        self.serialized = serialized
        self.response_schema = response_schema


def build_query_data_request(schema: 'pa.Schema' = pa.schema([]), filters: dict = None, field_names: list = None):
    filters = filters or {}

    builder = flatbuffers.Builder(1024)

    source_name = builder.CreateString('')  # required

    fields = [build_field(builder, f, f.name) for f in schema]

    fb_schema.StartFieldsVector(builder, len(fields))
    for offset in reversed(fields):
        builder.PrependUOffsetTRelative(offset)
    fields = builder.EndVector()

    fb_schema.Start(builder)
    fb_schema.AddFields(builder, fields)
    schema_obj = fb_schema.End(builder)

    predicate = Predicate(schema, filters)
    filter_obj = predicate.serialize(builder)

    parser = QueryDataParser(schema)
    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]

    response_schema = pa.schema([fields_map[name] for name in field_names])
    projection_fields = []
    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]
        _logger.info("projecting field=%s positions=%s", field_name, positions)
        for leaf_position in positions:
            fb_field_index.Start(builder)
            fb_field_index.AddPosition(builder, leaf_position)
            offset = fb_field_index.End(builder)
            projection_fields.append(offset)
    fb_source.StartProjectionVector(builder, len(projection_fields))
    for offset in reversed(projection_fields):
        builder.PrependUOffsetTRelative(offset)
    projection = builder.EndVector()

    fb_source.Start(builder)
    fb_source.AddName(builder, source_name)
    fb_source.AddSchema(builder, schema_obj)
    fb_source.AddFilter(builder, filter_obj)
    fb_source.AddProjection(builder, projection)
    source = fb_source.End(builder)

    fb_relation.Start(builder)
    fb_relation.AddImplType(builder, rel_impl.RelationImpl.Source)
    fb_relation.AddImpl(builder, source)
    relation = fb_relation.End(builder)

    builder.Finish(relation)
    return QueryDataRequest(serialized=builder.Output(), response_schema=response_schema)


def convert_column_types(table: 'pa.Table') -> 'pa.Table':
    """
    Adjusting table values

    1. Because the timestamp resolution is too high it is necessary to trim it. ORION-96961
    2. Since the values of nfs_mode_bits are returned in decimal, need to convert them to octal,
    as in all representations, so that the mode of 448 turn into 700
    3. for owner_name and group_owner_name 0 -> root, and 65534 -> nobody
    """
    ts_indexes = []
    indexes_of_fields_to_change = {}
    sid_to_name = {
        '0': 'root',
        '65534': 'nobody'  # NFSNOBODY_UID_16_BIT
    }
    column_matcher = {  # column_name: custom converting rule
        'nfs_mode_bits': lambda val: int(oct(val).replace('0o', '')) if val is not None else val,
        'owner_name': lambda val: sid_to_name.get(val, val),
        'group_owner_name': lambda val: sid_to_name.get(val, val),
    }
    for index, field in enumerate(table.schema):
        if isinstance(field.type, pa.TimestampType) and field.type.unit == 'ns':
            ts_indexes.append(index)
        if field.name in column_matcher:
            indexes_of_fields_to_change[field.name] = index
    for changing_index in ts_indexes:
        field_name = table.schema[changing_index].name
        _logger.info(f'changing resolution for {field_name} to us')
        new_column = table[field_name].cast(pa.timestamp('us'), safe=False)
        table = table.set_column(changing_index, field_name, new_column)
    for field_name, changing_index in indexes_of_fields_to_change.items():
        _logger.info(f'applying custom rules to {field_name}')
        new_column = table[field_name].to_pylist()
        new_column = list(map(column_matcher[field_name], new_column))
        new_column = pa.array(new_column, table[field_name].type)
        table = table.set_column(changing_index, field_name, new_column)
    return table
