vastdb 1.4.0__py3-none-any.whl → 2.0.1__py3-none-any.whl

vastdb/_ibis_support.py

@@ -0,0 +1,28 @@
+import ibis
+import pyarrow as pa
+from ibis.expr.types.structs import IbisError
+
+from vastdb import errors
+
+
+def validate_ibis_support_schema(arrow_schema: pa.Schema):
+    """Validate that the provided Arrow schema is compatible with Ibis.
+
+    Raises NotSupportedSchema if the schema contains unsupported fields.
+    """
+    unsupported_fields = []
+    first_exception = None
+    for f in arrow_schema:
+        try:
+            ibis.Schema.from_pyarrow(pa.schema([f]))
+        except (IbisError, ValueError, KeyError) as e:
+            if first_exception is None:
+                first_exception = e
+            unsupported_fields.append(f)
+
+    if unsupported_fields:
+        raise errors.NotSupportedSchema(
+            message=f"Ibis does not support the schema {unsupported_fields=}",
+            schema=arrow_schema,
+            cause=first_exception
+        )

vastdb/_internal.py

@@ -7,7 +7,7 @@ import time
 import urllib.parse
 from collections import defaultdict, namedtuple
 from enum import Enum
-from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Union, cast
 
 import backoff
 import flatbuffers
@@ -17,6 +17,7 @@ import requests
 import urllib3
 import xmltodict
 from aws_requests_auth.aws_auth import AWSRequestsAuth
+from ibis.expr.operations import Field, Node
 from ibis.expr.operations.generic import (
     IsNull,
     Literal,
@@ -34,7 +35,6 @@ from ibis.expr.operations.logical import (
     NotEquals,
     Or,
 )
-from ibis.expr.operations.relations import Field
 from ibis.expr.operations.strings import StartsWith, StringContains
 from ibis.expr.operations.structs import StructField
 
@@ -52,6 +52,7 @@ import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.Int8Literal as fb_
 import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.Int16Literal as fb_int16_lit
 import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.Int32Literal as fb_int32_lit
 import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.Int64Literal as fb_int64_lit
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.ListLiteral as fb_list_lit
 import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.Literal as fb_literal
 import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.Relation as fb_relation
 import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.RelationImpl as rel_impl
@@ -262,20 +263,32 @@ class Predicate:
                         node = nodes_map[name]
                         nodes_map = node.children_map
 
-                    # TODO: support predicate pushdown for leaf nodes (ORION-160338)
+                    literal_field = node.field
+                    literal_column_index = node.index
                     if node.children:
-                        raise NotImplementedError(node.field)  # no predicate pushdown for nested columns
-                    column_offset = self.build_column(position=node.index)
+                        # Support fixed size list of a single flat child.
+                        if pa.types.is_fixed_size_list(node.type) and len(node.children) == 1 and not node.children[
+                            0].children:
+                            # Similar to projection, the index of the column should be the leaf which is the child's.
+                            literal_column_index = node.children[0].index
+                            # Set the literal type to be a list rather than fixed list since fixed list is not supported.
+                            # https://github.com/apache/arrow/blob/apache-arrow-7.0.0/cpp/src/arrow/compute/exec/ir_consumer.cc#L287
+                            literal_field = node.field.with_type(pa.list_(node.field.type.value_field))
+                        else:
+                            # TODO: support predicate pushdown for leaf nodes (ORION-160338)
+                            raise NotImplementedError(node.field)  # no predicate pushdown for nested columns
+
+                    column_offset = self.build_column(position=literal_column_index)
                     for literal in literals:
                         args_offsets = [column_offset]
                         if literal is not None:
-                            args_offsets.append(self.build_literal(field=node.field, value=literal.value))
+                            args_offsets.append(self.build_literal_expression(field=literal_field, value=literal.value))
                         if builder_func == self.build_between:
-                            args_offsets.append(self.build_literal(field=node.field, value=lower.value))
-                            args_offsets.append(self.build_literal(field=node.field, value=upper.value))
+                            args_offsets.append(self.build_literal_expression(field=literal_field, value=lower.value))
+                            args_offsets.append(self.build_literal_expression(field=literal_field, value=upper.value))
                         if builder_func == self.build_starts_with:
-                            args_offsets.append(self.build_literal(field=node.field, value=lower_bytes))
-                            args_offsets.append(self.build_literal(field=node.field, value=upper_bytes))
+                            args_offsets.append(self.build_literal_expression(field=literal_field, value=lower_bytes))
+                            args_offsets.append(self.build_literal_expression(field=literal_field, value=upper_bytes))
 
                         inner_offsets.append(builder_func(*args_offsets))
 
@@ -326,14 +339,14 @@ class Predicate:
         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)
+                literal = self.build_literal_expression(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)
+            literal = self.build_literal_expression(field=field, value=value)
             rules.append(op(column, literal))
 
         return self.build_and(rules)
@@ -359,145 +372,93 @@ class Predicate:
     # see https://github.com/apache/arrow/blob/main/format/Schema.fbs
     # https://github.com/apache/arrow/blob/apache-arrow-7.0.0/experimental/computeir/Expression.fbs
     # https://github.com/apache/arrow/blob/apache-arrow-7.0.0/experimental/computeir/Literal.fbs
-    def build_literal(self, field: pa.Field, value):
-        literal_type: Any
-
-        if field.type.equals(pa.int64()) or field.type.equals(pa.uint64()):
-            is_signed = field.type.equals(pa.int64())
-            literal_type = fb_int64_lit if is_signed else fb_uint64_lit
-            literal_impl = LiteralImpl.Int64Literal if is_signed else LiteralImpl.UInt64Literal
-
-            field_type_type = Type.Int
-            fb_int.Start(self.builder)
-            fb_int.AddBitWidth(self.builder, field.type.bit_width)
-            fb_int.AddIsSigned(self.builder, is_signed)
-            field_type = fb_int.End(self.builder)
-
+    def build_literal_impl(self, pa_type: pa.DataType, value) -> Tuple[int, int]:
+        '''
+        Builds a LiteralImpl for the given Arrow type and value.
+        :param pa_type: Literal type as defined in Arrow.
+        :param value: Value to be used in the LiteralImpl.
+        :return: Tuple[LiteralImpl, buffer_value]
+        '''
+        if pa.types.is_integer(pa_type):
+            impl_type, impl_class = None, None
             value = int(value)
-        elif field.type.equals(pa.int32()) or field.type.equals(pa.uint32()):
-            is_signed = field.type.equals(pa.int32())
-            literal_type = fb_int32_lit if is_signed else fb_uint32_lit
-            literal_impl = LiteralImpl.Int32Literal if is_signed else LiteralImpl.UInt32Literal
 
-            field_type_type = Type.Int
-            fb_int.Start(self.builder)
-            fb_int.AddBitWidth(self.builder, field.type.bit_width)
-            fb_int.AddIsSigned(self.builder, is_signed)
-            field_type = fb_int.End(self.builder)
+            if pa.types.is_int8(pa_type):
+                impl_type, impl_class = LiteralImpl.Int8Literal, fb_int8_lit
+            elif pa.types.is_uint8(pa_type):
+                impl_type, impl_class = LiteralImpl.UInt8Literal, fb_uint8_lit
+            elif pa.types.is_int16(pa_type):
+                impl_type, impl_class = LiteralImpl.Int16Literal, fb_int16_lit
+            elif pa.types.is_uint16(pa_type):
+                impl_type, impl_class = LiteralImpl.UInt16Literal, fb_uint16_lit
+            elif pa.types.is_int32(pa_type):
+                impl_type, impl_class = LiteralImpl.Int32Literal, fb_int32_lit
+            elif pa.types.is_uint32(pa_type):
+                impl_type, impl_class = LiteralImpl.UInt32Literal, fb_uint32_lit
+            elif pa.types.is_int64(pa_type):
+                impl_type, impl_class = LiteralImpl.Int64Literal, fb_int64_lit
+            elif pa.types.is_uint64(pa_type):
+                impl_type, impl_class = LiteralImpl.UInt64Literal, fb_uint64_lit
+            else:
+                raise ValueError(f'unsupported integer predicate type: {pa_type}, value={value}')
 
-            value = int(value)
-        elif field.type.equals(pa.int16()) or field.type.equals(pa.uint16()):
-            is_signed = field.type.equals(pa.int16())
-            literal_type = fb_int16_lit if is_signed else fb_uint16_lit
-            literal_impl = LiteralImpl.Int16Literal if is_signed else LiteralImpl.UInt16Literal
+            impl_class.Start(self.builder)
+            impl_class.AddValue(self.builder, value)
+            buffer_value = impl_class.End(self.builder)
+            return impl_type, buffer_value
 
-            field_type_type = Type.Int
-            fb_int.Start(self.builder)
-            fb_int.AddBitWidth(self.builder, field.type.bit_width)
-            fb_int.AddIsSigned(self.builder, is_signed)
-            field_type = fb_int.End(self.builder)
+        if pa.types.is_floating(pa_type):
+            impl_type, impl_class = None, None
+            value = float(value)
 
-            value = int(value)
-        elif field.type.equals(pa.int8()) or field.type.equals(pa.uint8()):
-            is_signed = field.type.equals(pa.int8())
-            literal_type = fb_int8_lit if is_signed else fb_uint8_lit
-            literal_impl = LiteralImpl.Int8Literal if is_signed else LiteralImpl.UInt8Literal
+            if pa.types.is_float32(pa_type):
+                impl_type, impl_class = LiteralImpl.Float32Literal, fb_float32_lit
+            elif pa.types.is_float64(pa_type):
+                impl_type, impl_class = LiteralImpl.Float64Literal, fb_float64_lit
+            else:
+                # Float16 is not supported by Vast.
+                raise ValueError(f'unsupported floating point predicate type: {pa_type}, value={value}')
 
-            field_type_type = Type.Int
-            fb_int.Start(self.builder)
-            fb_int.AddBitWidth(self.builder, field.type.bit_width)
-            fb_int.AddIsSigned(self.builder, is_signed)
-            field_type = fb_int.End(self.builder)
+            impl_class.Start(self.builder)
+            impl_class.AddValue(self.builder, value)
+            buffer_value = impl_class.End(self.builder)
+            return impl_type, buffer_value
 
-            value = int(value)
-        elif field.type.equals(pa.float32()):
-            literal_type = fb_float32_lit
-            literal_impl = LiteralImpl.Float32Literal
+        if pa_type.equals(pa.string()):
+            value = self.builder.CreateString(value)
 
-            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)
+            fb_string_lit.Start(self.builder)
+            fb_string_lit.AddValue(self.builder, value)
+            buffer_value = fb_string_lit.End(self.builder)
+            return LiteralImpl.StringLiteral, buffer_value
 
-            value = float(value)
-        elif field.type.equals(pa.float64()):
-            literal_type = fb_float64_lit
-            literal_impl = LiteralImpl.Float64Literal
+        if pa_type.equals(pa.date32()):  # pa.date64() is not supported
+            # Assuming units are in Days. Look at get_field_type for more details.
+            value, = pa.array([value], pa_type).cast(pa.int32()).to_pylist()
 
-            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)
+            fb_date32_lit.Start(self.builder)
+            fb_date32_lit.AddValue(self.builder, value)
+            buffer_value = fb_date32_lit.End(self.builder)
+            return LiteralImpl.DateLiteral, buffer_value
 
-            value = float(value)
-        elif field.type.equals(pa.string()):
-            literal_type = fb_string_lit
-            literal_impl = LiteralImpl.StringLiteral
+        if pa.types.is_timestamp(pa_type):
+            value, = pa.array([value], pa_type).cast(pa.int64()).to_pylist()
 
-            field_type_type = Type.Utf8
-            fb_utf8.Start(self.builder)
-            field_type = fb_utf8.End(self.builder)
+            fb_timestamp_lit.Start(self.builder)
+            fb_timestamp_lit.AddValue(self.builder, value)
+            buffer_value = fb_timestamp_lit.End(self.builder)
+            return LiteralImpl.TimestampLiteral, buffer_value
 
-            value = self.builder.CreateString(value)
-        elif field.type.equals(pa.date32()):  # pa.date64() is not supported
-            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)
-            value, = pa.array([value], field.type).cast(pa.int32()).to_pylist()
-        elif isinstance(field.type, pa.TimestampType):
-            literal_type = fb_timestamp_lit
-            literal_impl = LiteralImpl.TimestampLiteral
-
-            if field.type.equals(pa.timestamp('s')):
-                unit = TimeUnit.SECOND
-            if field.type.equals(pa.timestamp('ms')):
-                unit = TimeUnit.MILLISECOND
-            if field.type.equals(pa.timestamp('us')):
-                unit = TimeUnit.MICROSECOND
-            if field.type.equals(pa.timestamp('ns')):
-                unit = TimeUnit.NANOSECOND
-
-            field_type_type = Type.Timestamp
-            fb_timestamp.Start(self.builder)
-            fb_timestamp.AddUnit(self.builder, unit)
-            field_type = fb_timestamp.End(self.builder)
-            value, = pa.array([value], field.type).cast(pa.int64()).to_pylist()
-        elif isinstance(field.type, (pa.Time32Type, pa.Time64Type)):
-            literal_type = fb_time_lit
-            literal_impl = LiteralImpl.TimeLiteral
-
-            if field.type.equals(pa.time32('s')):
-                target_type = pa.int32()
-                unit = TimeUnit.SECOND
-            if field.type.equals(pa.time32('ms')):
-                target_type = pa.int32()
-                unit = TimeUnit.MILLISECOND
-            if field.type.equals(pa.time64('us')):
-                target_type = pa.int64()
-                unit = TimeUnit.MICROSECOND
-            if field.type.equals(pa.time64('ns')):
-                target_type = pa.int64()
-                unit = TimeUnit.NANOSECOND
-
-            field_type_type = Type.Time
-            fb_time.Start(self.builder)
-            fb_time.AddBitWidth(self.builder, field.type.bit_width)
-            fb_time.AddUnit(self.builder, unit)
-            field_type = fb_time.End(self.builder)
-
-            value, = pa.array([value], field.type).cast(target_type).to_pylist()
-        elif field.type.equals(pa.bool_()):
-            literal_type = fb_bool_lit
-            literal_impl = LiteralImpl.BooleanLiteral
-
-            field_type_type = Type.Bool
-            fb_bool.Start(self.builder)
-            field_type = fb_bool.End(self.builder)
+        if pa.types.is_time(pa_type):
+            target_type = pa.int32() if pa.types.is_time32(pa_type) else pa.int64()
+            value, = pa.array([value], pa_type).cast(target_type).to_pylist()
 
+            fb_time_lit.Start(self.builder)
+            fb_time_lit.AddValue(self.builder, value)
+            buffer_value = fb_time_lit.End(self.builder)
+            return LiteralImpl.TimeLiteral, buffer_value
+
+        if pa_type.equals(pa.bool_()):
             # Handle both boolean values and string representations
             if isinstance(value, bool):
                 value = value
@@ -505,46 +466,65 @@ class Predicate:
                 value = value.lower() == 'true'
             else:
                 value = bool(value)
-        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')
 
+            fb_bool_lit.Start(self.builder)
+            fb_bool_lit.AddValue(self.builder, value)
+            buffer_value = fb_bool_lit.End(self.builder)
+            return LiteralImpl.BooleanLiteral, buffer_value
+
+        if pa.types.is_decimal128(pa_type):
+            int_value = int(float(value) * 10 ** pa_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)
+            fb_decimal_lit.Start(self.builder)
+            fb_decimal_lit.AddValue(self.builder, value)
+            buffer_value = fb_decimal_lit.End(self.builder)
+            return LiteralImpl.DecimalLiteral, buffer_value
 
+        if pa_type.equals(pa.binary()):
             value = self.builder.CreateByteVector(value)
-        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_binary_lit.Start(self.builder)
+            fb_binary_lit.AddValue(self.builder, value)
+            buffer_value = fb_binary_lit.End(self.builder)
+            return LiteralImpl.BinaryLiteral, buffer_value
+
+        # pa.types.is_list is False for FixedSizeList which is important since parsing of FixedSizeList is not supported
+        # https://github.com/apache/arrow/blob/apache-arrow-7.0.0/cpp/src/arrow/compute/exec/ir_consumer.cc#L287
+        if pa.types.is_list(pa_type):
+            pa_type = cast(pa.FixedSizeListType, pa_type)
+
+            buffer_literals = []
+            for element in value:
+                buffer_literals.append(self.build_literal(pa_type.value_field, element))
+            fb_list_lit.StartValuesVector(self.builder, len(buffer_literals))
+            for offset in reversed(buffer_literals):
+                self.builder.PrependUOffsetTRelative(offset)
+            values_buffer = self.builder.EndVector()
+
+            fb_list_lit.Start(self.builder)
+            fb_list_lit.AddValues(self.builder, values_buffer)
+            buffer_value = fb_list_lit.End(self.builder)
+            return LiteralImpl.ListLiteral, buffer_value
+
+        raise ValueError(f'unsupported literal type={pa_type}, value={value}')
 
-        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)
+    def build_literal(self, field: pa.Field, value) -> int:
+        literal_impl_type, literal_impl_buffer = self.build_literal_impl(field.type, value)
+
+        # Literal type should not contain name for more information
+        # https://github.com/apache/arrow/blob/apache-arrow-7.0.0/cpp/src/arrow/compute/exec/ir_consumer.cc#L326
+        field_buffer = build_field(self.builder, field, include_name=False)
 
         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_literal.AddImplType(self.builder, literal_impl_type)
+        fb_literal.AddImpl(self.builder, literal_impl_buffer)
+        fb_literal.AddType(self.builder, field_buffer)
+        return fb_literal.End(self.builder)
 
+    def build_literal_expression(self, field: pa.Field, value) -> int:
+        buffer_literal = self.build_literal(field, value)
         fb_expression.Start(self.builder)
         fb_expression.AddImplType(self.builder, ExpressionImpl.Literal)
         fb_expression.AddImpl(self.builder, buffer_literal)
@@ -937,7 +917,7 @@ class VastdbApi:
         """Make sure that the connections closed."""
         self._session.close()
 
-    def with_endpoint(self, endpoint):
+    def with_endpoint(self, endpoint) -> 'VastdbApi':
         """Open a new session for targeting a specific endpoint."""
         return VastdbApi(endpoint=endpoint,
             access_key=self.access_key,
@@ -1612,7 +1592,7 @@ class VastdbApi:
 
         return headers
 
-    def _build_query_data_url_params(self, projection, query_imports_table):
+    def _build_query_data_url_params(self, projection: Optional[str], query_imports_table):
         if query_imports_table and projection:
             raise ValueError("Can't query both imports and projection table")
 
@@ -1624,8 +1604,8 @@ class VastdbApi:
         return url_params
 
     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,
+                   txid: Optional[int] = 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: Optional[str] = '', enable_sorted_projections=True,
                    request_format='string', response_format='string', query_imports_table=False):
         """
         GET /mybucket/myschema/mytable?data HTTP/1.1
@@ -2305,16 +2285,16 @@ def get_field_type(builder: flatbuffers.Builder, field: pa.Field):
     return field_type, field_type_type
 
 
-def build_field(builder: flatbuffers.Builder, f: pa.Field, name: str):
+def build_field(builder: flatbuffers.Builder, f: pa.Field, include_name=True):
     children = None
     if isinstance(f.type, pa.StructType):
-        children = [build_field(builder, child, child.name) for child in list(f.type)]
+        children = [build_field(builder, child, include_name) for child in list(f.type)]
     if pa.types.is_list(f.type) or pa.types.is_fixed_size_list(f.type):
-        children = [build_field(builder, f.type.value_field, "item")]
+        children = [build_field(builder, f.type.value_field.with_name("item"), include_name)]
     if isinstance(f.type, pa.MapType):
         children = [
-            build_field(builder, f.type.key_field, "key"),
-            build_field(builder, f.type.item_field, "value"),
+            build_field(builder, f.type.key_field.with_name("key"), include_name),
+            build_field(builder, f.type.item_field.with_name("value"), include_name),
         ]
 
         # adding "entries" column:
@@ -2340,10 +2320,15 @@ def build_field(builder: flatbuffers.Builder, f: pa.Field, name: str):
             builder.PrependUOffsetTRelative(offset)
         children = builder.EndVector()
 
-    col_name = builder.CreateString(name)
+    col_name = None
+    if include_name:
+        col_name = builder.CreateString(f.name)
+
     field_type, field_type_type = get_field_type(builder, f)
     fb_field.Start(builder)
-    fb_field.AddName(builder, col_name)
+    if col_name is not None:
+        fb_field.AddName(builder, col_name)
+    fb_field.AddNullable(builder, f.nullable)
     fb_field.AddTypeType(builder, field_type_type)
     fb_field.AddType(builder, field_type)
     if children is not None:
@@ -2358,19 +2343,48 @@ class QueryDataRequest:
         self.response_parser = response_parser
 
 
-def get_response_schema(schema: 'pa.Schema' = pa.schema([]), field_names: Optional[List[str]] = None):
+def get_response_schema(schema: 'pa.Schema' = pa.schema([]), field_names: Optional[List[str]] = None) -> pa.Schema:
     if field_names is None:
         field_names = [field.name for field in schema]
 
     return pa.schema([schema.field(name) for name in field_names])
 
 
-def build_query_data_request(schema: 'pa.Schema' = pa.schema([]), predicate: ibis.expr.types.BooleanColumn = None, field_names: Optional[List[str]] = None):
+def _column_names_in_node_tree(expr: ibis.expr.types.Expr) -> set[str]:
+    def walk_op(op: Node):
+        if isinstance(op, Field):
+            names.add(op.name)
+        elif isinstance(op, (list, tuple)):
+            for item in op:
+                walk_op(item)
+        else:
+            for arg in getattr(op, "args", ()):
+                walk_op(arg)
+
+    names: set[str] = set()
+    walk_op(expr.op())
+    return names
+
+
+def build_query_data_request(schema: pa.Schema,
+                             predicate: ibis.expr.types.BooleanColumn = None,
+                             field_names: Optional[list[str]] = None) -> QueryDataRequest:
+    if field_names is None:
+        queried_columns = [f.name for f in schema]
+    else:
+        # apparently there are some tests that send a tuple despite the signature asking for a list
+        queried_columns = list(field_names)
+
+    column_names_in_predicate = set() if predicate is None else _column_names_in_node_tree(predicate)
+    column_names_required_by_predicate = column_names_in_predicate - set(queried_columns)
+    queried_columns.extend(column_names_required_by_predicate)
+    schema = pa.schema((schema.field(name) for name in queried_columns))
+
     builder = flatbuffers.Builder(1024)
 
     source_name = builder.CreateString('')  # required
 
-    fields = [build_field(builder, f, f.name) for f in schema]
+    fields = [build_field(builder, f) for f in schema]
 
     fb_schema.StartFieldsVector(builder, len(fields))
     for offset in reversed(fields):
@@ -2387,18 +2401,17 @@ def build_query_data_request(schema: 'pa.Schema' = pa.schema([]), predicate: ibi
     parser = QueryDataParser(schema)
     leaves_map = {node.field.name: [leaf.index for leaf in node._iter_leaves()] for node in parser.nodes}
 
-    response_schema = get_response_schema(schema, field_names)
-    field_names = [field.name for field in response_schema]
-
     projection_fields = []
-    for field_name in field_names:
+    for field in schema:
         # TODO: only root-level projection pushdown is supported (i.e. no support for SELECT s.x FROM t)
-        positions = leaves_map[field_name]
+        positions = leaves_map[field.name]
+
         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)
@@ -2418,4 +2431,4 @@ def build_query_data_request(schema: 'pa.Schema' = pa.schema([]), predicate: ibi
 
     builder.Finish(relation)
 
-    return QueryDataRequest(serialized=builder.Output(), response_schema=response_schema, response_parser=QueryDataParser(response_schema))
+    return QueryDataRequest(serialized=builder.Output(), response_schema=schema, response_parser=QueryDataParser(schema))