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

vastdb/_internal.py

@@ -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
 
@@ -2350,9 +2350,36 @@ def get_response_schema(schema: 'pa.Schema' = pa.schema([]), field_names: Option
     return pa.schema([schema.field(name) for name in field_names])
 
 
-def build_query_data_request(schema: 'pa.Schema' = pa.schema([]),
+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:
+                             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
@@ -2374,18 +2401,17 @@ def build_query_data_request(schema: 'pa.Schema' = pa.schema([]),
     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)
@@ -2405,4 +2431,4 @@ def build_query_data_request(schema: 'pa.Schema' = pa.schema([]),
 
     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))

{vastdb-2.0.0.dist-info → vastdb-2.0.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vastdb
-Version: 2.0.0
+Version: 2.0.1
 Summary: VAST Data SDK
 Home-page: https://github.com/vast-data/vastdb_sdk
 Author: VAST DATA

{vastdb-2.0.0.dist-info → vastdb-2.0.1.dist-info}/RECORD

@@ -1,6 +1,6 @@
 vastdb/__init__.py,sha256=uf-AXdzsD4nPxFP7WxkcAXGG0whv8BHLrrXCJtsPGaQ,436
 vastdb/_ibis_support.py,sha256=sJieOMvDWpsciPKh1mJzS56jxLtCRVlvK41hW84vexM,866
-vastdb/_internal.py,sha256=ZR9da7t-dQGtax0Cz-3CgtX1C0nLKc5VI9dkfO349RY,108764
+vastdb/_internal.py,sha256=fkMCzhPPFFC6-EmeJz9R8Krjbi_u9UXgv_HLvhaigHo,109583
 vastdb/_table_interface.py,sha256=EdS50_x9ZVzRWDyV6b5hB9qIUdA7S4hNBgNq3t8deHo,3611
 vastdb/bucket.py,sha256=ulkTI_7xw5FYVzcrTFC7G2ijmTTVSmvJZUdgzycGHR0,2588
 vastdb/config.py,sha256=OehnsWrjzv0-SUouEXmkrKBugiWyhXOn4XiSLV3s9yk,2342
@@ -214,8 +214,8 @@ vastdb/vast_flatbuf/tabular/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMp
 vastdb/vast_tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vastdb/vast_tests/test_ha.py,sha256=744P4G6VJ09RIkHhMQL4wlipCBJWQVMhyvUrSc4k1HQ,975
 vastdb/vast_tests/test_scale.py,sha256=5jGwOdZH6Tv5tPdZYPWoqcxOceI2jA5i2D1zNKZHER4,3958
-vastdb-2.0.0.dist-info/licenses/LICENSE,sha256=obffan7LYrq7hLHNrY7vHcn2pKUTBUYXMKu-VOAvDxU,11333
-vastdb-2.0.0.dist-info/METADATA,sha256=GeJUYt9gomO12Gouliy_4TtHlx1-IwuZU6hwtpaT9mA,1686
-vastdb-2.0.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-vastdb-2.0.0.dist-info/top_level.txt,sha256=nnKAaZaQa8GFbYpWAexr_B9HrhonZbUlX6hL6AC--yA,7
-vastdb-2.0.0.dist-info/RECORD,,
+vastdb-2.0.1.dist-info/licenses/LICENSE,sha256=obffan7LYrq7hLHNrY7vHcn2pKUTBUYXMKu-VOAvDxU,11333
+vastdb-2.0.1.dist-info/METADATA,sha256=3AyWHxel0lNv7ORfAoUhZXnol-vuGM12aprOBLAVVKo,1686
+vastdb-2.0.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+vastdb-2.0.1.dist-info/top_level.txt,sha256=nnKAaZaQa8GFbYpWAexr_B9HrhonZbUlX6hL6AC--yA,7
+vastdb-2.0.1.dist-info/RECORD,,