Merge pull request #2701 from mabel-dev/microoptimizations/2

improve test suite performance

Author: joocer

.github/workflows/release.yaml

@@ -127,9 +127,6 @@ jobs:
       - name: Check out repository
         uses: actions/checkout@v4.1.7
 
-      - name: Install Python package dependencies
-        run: python -m pip install --upgrade cython==3.1.2 wheel numpy setuptools_rust pyarrow
-
       - name: Create source dist
         run: python setup.py sdist
 
@@ -232,6 +229,4 @@ jobs:
       - run: mv -v dist-win-3.12/* dist/
 
       - name: Publish distribution 📦 to PyPI
-        uses: pypa/gh-action-pypi-publish@release/v1
-        with:
-          password: ${{ secrets.PYPI_API_TOKEN }}
+        uses: pypa/gh-action-pypi-publish@release/v1

opteryx/__version__.py

@@ -1,4 +1,4 @@
-__build__ = 1393
+__build__ = 1394
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

opteryx/operators/read_node.py

@@ -43,10 +43,10 @@ def struct_to_jsonb(table: pyarrow.Table) -> pyarrow.Table:
         # Check if the column is a STRUCT
         if pyarrow.types.is_struct(field.type):
             # Convert each row in the STRUCT column to a JSON string
-            json_strings = [
-                orjson.dumps(row.as_py()) if row.is_valid else None for row in table.column(i)
-            ]
-            json_array = pyarrow.array(json_strings, type=pyarrow.binary())
+            json_array = pyarrow.array(
+                [None if row is None else orjson.dumps(row) for row in table.column(i).to_pylist()],
+                type=pyarrow.binary(),
+            )
 
             # Drop the original STRUCT column
             table = table.drop_columns(field.name)
@@ -62,13 +62,13 @@ def struct_to_jsonb(table: pyarrow.Table) -> pyarrow.Table:
 
             # Convert each list element
             converted_data = []
-            for item in list_array:
+            for item in list_array.to_pylist():
                 if item is None:
                     converted_data.append(None)
                 else:
                     # Each item is a list of structs
                     converted_list = []
-                    for struct in item.as_py():
+                    for struct in item:
                         if struct is None:
                             converted_list.append(None)
                         else:
@@ -98,28 +98,28 @@ def normalize_morsel(schema: RelationSchema, morsel: pyarrow.Table) -> pyarrow.T
     # rename columns for internal use
     target_column_names = []
     # columns in the data but not in the schema, droppable
-    droppable_columns = []
+    droppable_columns = set()
 
     # Find which columns to drop and which columns we already have
     for i, column in enumerate(morsel.column_names):
         column_name = schema.find_column(column)
         if column_name is None:
-            droppable_columns.append(i)
+            droppable_columns.add(i)
         else:
             target_column_names.append(str(column_name))
 
     # Remove from the end otherwise we'll remove the wrong columns after we've removed one
-    droppable_columns.reverse()
-    for droppable in droppable_columns:
-        morsel = morsel.remove_column(droppable)
+    if droppable_columns:
+        keep_indices = [i for i in range(len(morsel.columns)) if i not in droppable_columns]
+        morsel = morsel.select(keep_indices)
 
     # remane columns to the internal names (identities)
     morsel = morsel.rename_columns(target_column_names)
 
     # add columns we don't have, populate with nulls but try to get the correct type
     for column in schema.columns:
         if column.identity not in target_column_names:
-            null_column = pyarrow.array([None] * morsel.num_rows, type=column.arrow_field.type)
+            null_column = pyarrow.nulls(morsel.num_rows, type=column.arrow_field.type)
             field = pyarrow.field(name=column.identity, type=column.arrow_field.type)
             morsel = morsel.append_column(field, null_column)
 

opteryx/planner/binder/binder_visitor.py

@@ -4,6 +4,7 @@
 # Distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND.
 
 import re
+from functools import lru_cache
 from typing import List
 from typing import Set
 from typing import Tuple
@@ -216,6 +217,11 @@ def convert_using_to_on(
     return conditions[0]
 
 
+@lru_cache(maxsize=128)
+def node_type_to_method_name(node_type: str) -> str:
+    return f"visit_{CAMEL_TO_SNAKE.sub('_', node_type).lower()}"
+
+
 class BinderVisitor:
     """
     The BinderVisitor visits each node in the query plan and adds catalogue information
@@ -242,13 +248,13 @@ def visit_node(self, node: Node, context: BindingContext) -> Tuple[Node, Binding
             The node and context after binding.
         """
         node_type = node.node_type.name  # type:ignore
-        visit_method_name = f"visit_{CAMEL_TO_SNAKE.sub('_', node_type).lower()}"
+        visit_method_name = node_type_to_method_name(node_type)
         visit_method = getattr(self, visit_method_name, None)
         if visit_method is None:
             # DEBUG: print(f"BinderVisitor: No method found for {visit_method_name}")
             return node, context
 
-        return_node, return_context = visit_method(node.copy(), context.copy())
+        return_node, return_context = visit_method(node, context)
 
         # DEBUG: from opteryx.exceptions import InvalidInternalStateError
         # DEBUG:

opteryx/third_party/travers/graph.py

@@ -265,7 +265,12 @@ def outgoing_edges(self, source: str) -> List[Tuple]:
         Returns:
             Set of Tuples (Source, Target, Relationship)
         """
-        return [(source, t, r) for t, r in self._edges.get(source, tuple())]
+        try:
+            edges = self._edges[source]
+        except KeyError:
+            return []
+
+        return [(source, target, relation) for target, relation in edges]
 
     def ingoing_edges(self, target) -> List[Tuple]:
         """

opteryx/utils/file_decoders.py

@@ -87,10 +87,10 @@ def convert_arrow_schema_to_orso_schema(
 
 def get_decoder(dataset: str) -> Callable:
     """helper routine to get the decoder for a given file"""
-    ext = dataset.split(".")[-1].lower()
-    if ext not in KNOWN_EXTENSIONS:  # pragma: no cover
-        raise UnsupportedFileTypeError(f"Unsupported file type - {ext}")
-    file_decoder, file_type = KNOWN_EXTENSIONS[ext]
+    ext = dataset.rpartition(".")[2].lower()
+    file_decoder, file_type = KNOWN_EXTENSIONS.get(ext, (None, None))
+    if file_type is None:
+        raise UnsupportedFileTypeError(f"Unsupported file type: {ext}")
     if file_type != ExtentionType.DATA:  # pragma: no cover
         return do_nothing
     return file_decoder
@@ -387,12 +387,12 @@ def jsonl_decoder(
     from opteryx.third_party.tktech import csimdjson as simdjson
 
     if isinstance(buffer, memoryview):
-        buffer = MemoryViewStream(buffer)
+        # If it's a memoryview, we need to convert it to bytes
+        buffer = buffer.tobytes()
     if not isinstance(buffer, bytes):
         buffer = buffer.read()
 
     parser = simdjson.Parser()
-    lines = buffer.split(b"\n")
 
     # preallocate and reuse dicts
     rows = []
@@ -402,17 +402,27 @@ def jsonl_decoder(
         # If projection is specified, we only need to ensure we keep the projected keys
         keys_union = {c.value for c in projection}
 
-    for line in lines:
+    start = 0
+    end = len(buffer)
+
+    while start < end:
+        newline = buffer.find(b"\n", start)
+        if newline == -1:
+            newline = end
+        line = buffer[start:newline]
+        start = newline + 1
+
         if not line:
             continue
+
         record = parser.parse(line)
 
+        # convert nested objects to string
+        row = record.as_dict()
         # keep track of all keys for schema padding
         if not projection:
-            keys_union.update(record.keys())
+            keys_union.update(row.keys())
 
-        # convert nested objects to string
-        row = record.as_dict()
         for key in keys_union:
             if isinstance(row.get(key), dict):
                 row[key] = record[key].mini