introduce equi joins

Author: joocer

.gitignore

@@ -218,3 +218,4 @@ third_party/mabel/draken/core/ops.c*
 tests/draken/performance/.perf_baseline.json
 opteryx/draken/vectors/_hash_api.c*
 third_party/mabel/draken/vectors/interval_vector.c*
+third_party/mabel/draken/morsels/align.c*

examples/advanced/align_tables_examples.py

@@ -0,0 +1,221 @@
+"""
+Simple example demonstrating align_tables usage for join operations.
+"""
+import numpy as np
+import pyarrow as pa
+
+from opteryx.draken import Morsel
+from opteryx.draken.morsels import align_tables_pyarray
+
+
+def example_1_basic_alignment():
+    """Basic example: align two tables by indices."""
+    print("Example 1: Basic Alignment")
+    print("-" * 50)
+    
+    # Create left table
+    left = pa.table({
+        "user_id": pa.array([1, 2, 3, 4, 5], type=pa.int64()),
+        "name": pa.array(["Alice", "Bob", "Charlie", "David", "Eve"]),
+        "age": pa.array([25, 30, 35, 40, 45], type=pa.int64()),
+    })
+    
+    # Create right table
+    right = pa.table({
+        "order_id": pa.array([101, 102, 103, 104, 105], type=pa.int64()),
+        "amount": pa.array([100.0, 200.0, 150.0, 300.0, 250.0]),
+    })
+    
+    # Convert to morsels
+    left_morsel = Morsel.from_arrow(left)
+    right_morsel = Morsel.from_arrow(right)
+    
+    # Align: take rows [0, 2, 4] from left and [1, 3, 4] from right
+    left_indices = np.array([0, 2, 4], dtype=np.int32)
+    right_indices = np.array([1, 3, 4], dtype=np.int32)
+    
+    result = align_tables_pyarray(left_morsel, right_morsel, left_indices, right_indices)
+    
+    print(f"Input shapes: {left_morsel.shape} + {right_morsel.shape}")
+    print(f"Output shape: {result.shape}")
+    print(f"Columns: {result.column_names}\n")
+    
+    result_arrow = result.to_arrow()
+    print(result_arrow)
+    print()
+
+
+def example_2_inner_join():
+    """Example 2: Simulate an inner join."""
+    print("Example 2: Inner Join Simulation")
+    print("-" * 50)
+    
+    # Users table
+    users = pa.table({
+        "user_id": pa.array([1, 2, 3, 4, 5], type=pa.int64()),
+        "name": pa.array(["Alice", "Bob", "Charlie", "David", "Eve"]),
+    })
+    
+    # Orders table
+    orders = pa.table({
+        "order_id": pa.array([101, 102, 103, 104], type=pa.int64()),
+        "user_id": pa.array([2, 1, 4, 2], type=pa.int64()),  # Foreign key
+        "amount": pa.array([100.0, 200.0, 150.0, 300.0]),
+    })
+    
+    users_morsel = Morsel.from_arrow(users)
+    orders_morsel = Morsel.from_arrow(orders)
+    
+    # Simulate hash join: match user_id
+    # orders.user_id = [2, 1, 4, 2]
+    # Match to users indices: [1, 0, 3, 1]
+    user_indices = np.array([1, 0, 3, 1], dtype=np.int32)  # Bob, Alice, David, Bob
+    order_indices = np.array([0, 1, 2, 3], dtype=np.int32)  # All orders
+    
+    result = align_tables_pyarray(users_morsel, orders_morsel, user_indices, order_indices)
+    
+    print("Joined users with their orders:")
+    print(result.to_arrow())
+    print()
+
+
+def example_3_duplicate_columns():
+    """Example 3: Handle duplicate column names."""
+    print("Example 3: Duplicate Column Handling")
+    print("-" * 50)
+    
+    left = pa.table({
+        "id": pa.array([1, 2, 3], type=pa.int64()),
+        "value": pa.array([10, 20, 30], type=pa.int64()),
+    })
+    
+    right = pa.table({
+        "id": pa.array([4, 5, 6], type=pa.int64()),  # Same column name!
+        "score": pa.array([100, 200, 300], type=pa.int64()),
+    })
+    
+    left_morsel = Morsel.from_arrow(left)
+    right_morsel = Morsel.from_arrow(right)
+    
+    indices = np.array([0, 1, 2], dtype=np.int32)
+    
+    result = align_tables_pyarray(left_morsel, right_morsel, indices, indices)
+    
+    print("Note: 'id' from right table is ignored (left takes precedence)")
+    print(f"Result columns: {result.column_names}")
+    print(result.to_arrow())
+    print()
+
+
+def example_4_large_scale():
+    """Example 4: Large-scale performance."""
+    print("Example 4: Large-Scale Performance")
+    print("-" * 50)
+    
+    import time
+    
+    n_rows = 1_000_000
+    
+    # Large left table
+    left = pa.table({
+        "id": pa.array(np.arange(n_rows, dtype=np.int64)),
+        "val1": pa.array(np.random.rand(n_rows)),
+        "val2": pa.array(np.random.rand(n_rows)),
+        "val3": pa.array(np.random.rand(n_rows)),
+    })
+    
+    # Large right table
+    right = pa.table({
+        "metric1": pa.array(np.random.rand(n_rows)),
+        "metric2": pa.array(np.random.rand(n_rows)),
+    })
+    
+    left_morsel = Morsel.from_arrow(left)
+    right_morsel = Morsel.from_arrow(right)
+    
+    # Select 50% of rows
+    sample_size = n_rows // 2
+    indices = np.random.choice(n_rows, size=sample_size, replace=True).astype(np.int32)
+    
+    print(f"Left table: {left_morsel.shape}")
+    print(f"Right table: {right_morsel.shape}")
+    print(f"Sample size: {sample_size:,} rows")
+    
+    # Warmup
+    _ = align_tables_pyarray(left_morsel, right_morsel, indices, indices)
+    
+    # Benchmark
+    num_runs = 5
+    times = []
+    
+    for _ in range(num_runs):
+        start = time.perf_counter()
+        result = align_tables_pyarray(left_morsel, right_morsel, indices, indices)
+        elapsed = time.perf_counter() - start
+        times.append(elapsed)
+    
+    avg_time = np.mean(times) * 1000
+    throughput = sample_size / (avg_time / 1000) / 1_000_000
+    
+    print("\nPerformance:")
+    print(f"  Average time: {avg_time:.1f} ms")
+    print(f"  Throughput:   {throughput:.1f} M rows/sec")
+    print(f"  Result shape: {result.shape}")
+    print()
+
+
+def example_5_mixed_types():
+    """Example 5: Mixed data types."""
+    print("Example 5: Mixed Data Types")
+    print("-" * 50)
+    
+    left = pa.table({
+        "int_col": pa.array([1, 2, 3], type=pa.int64()),
+        "float_col": pa.array([1.1, 2.2, 3.3]),
+        "string_col": pa.array(["a", "b", "c"]),
+        "bool_col": pa.array([True, False, True]),
+    })
+    
+    right = pa.table({
+        "date_col": pa.array([
+            pa.scalar(1, type=pa.date32()),
+            pa.scalar(2, type=pa.date32()),
+            pa.scalar(3, type=pa.date32()),
+        ]),
+        "timestamp_col": pa.array([
+            pa.scalar(1000, type=pa.timestamp('us')),
+            pa.scalar(2000, type=pa.timestamp('us')),
+            pa.scalar(3000, type=pa.timestamp('us')),
+        ]),
+    })
+    
+    left_morsel = Morsel.from_arrow(left)
+    right_morsel = Morsel.from_arrow(right)
+    
+    indices = np.array([0, 2], dtype=np.int32)
+    
+    result = align_tables_pyarray(left_morsel, right_morsel, indices, indices)
+    
+    print(f"Handling {result.num_columns} columns of different types:")
+    for name, dtype in zip(result.column_names, result.column_types):
+        print(f"  {name.decode('utf-8')}: {dtype}")
+    
+    print(f"\nResult:\n{result.to_arrow()}")
+    print()
+
+
+if __name__ == "__main__":
+    print("=" * 50)
+    print("align_tables Usage Examples")
+    print("=" * 50)
+    print()
+    
+    example_1_basic_alignment()
+    example_2_inner_join()
+    example_3_duplicate_columns()
+    example_4_large_scale()
+    example_5_mixed_types()
+    
+    print("=" * 50)
+    print("All examples completed!")
+    print("=" * 50)

examples/non_equi_join_example.py

@@ -6,10 +6,11 @@
 """
 
 import pyarrow as pa
+
+from opteryx import EOS
 from opteryx.compiled.joins import non_equi_nested_loop_join
-from opteryx.operators import NonEquiJoinNode
 from opteryx.models import QueryProperties
-from opteryx import EOS
+from opteryx.operators import NonEquiJoinNode
 
 
 def example_basic_non_equi_join():

opteryx/__version__.py

@@ -1,9 +1,9 @@
 # THIS FILE IS AUTOMATICALLY UPDATED DURING THE BUILD PROCESS
 # DO NOT EDIT THIS FILE DIRECTLY
 
-__build__ = 1830
+__build__ = 1836
 __author__ = "@joocer"
-__version__ = "0.26.2-beta.1830"
+__version__ = "0.26.2-beta.1836"
 
 # Store the version here so:
 # 1) we don't load dependencies by storing it in __init__.py

…eryx/compiled/joins/nested_loop_join.pyx …mpiled/joins/nested_loop_join_equals.pyxopteryx/compiled/joins/nested_loop_join.pyx renamed to opteryx/compiled/joins/nested_loop_join_equals.pyx opteryx/compiled/joins/nested_loop_join.pyx renamed to opteryx/compiled/joins/nested_loop_join_equals.pyx

@@ -15,18 +15,14 @@ for efficient columnar comparison operations.
 Supports: !=, >, >=, <, <=
 """
 
-import numpy
-cimport numpy
-numpy.import_array()
-
 from libc.stdint cimport int64_t
 
 from opteryx.compiled.structures.buffers cimport IntBuffer
-
+from opteryx.draken.morsels.morsel cimport Morsel
 
 cpdef tuple non_equi_nested_loop_join(
-    left_relation,
-    right_relation,
+    Morsel left_morsel,
+    Morsel right_morsel,
     str left_column,
     str right_column,
     str comparison_op
@@ -50,15 +46,11 @@ cpdef tuple non_equi_nested_loop_join(
     """
     from opteryx.draken.morsels.morsel import Morsel
 
-    # Convert Arrow tables to draken morsels
-    cdef object left_morsel = Morsel.from_arrow(left_relation)
-    cdef object right_morsel = Morsel.from_arrow(right_relation)
-
-    cdef int64_t left_rows = left_relation.num_rows
-    cdef int64_t right_rows = right_relation.num_rows
+    cdef int64_t left_rows = left_morsel.num_rows
+    cdef int64_t right_rows = right_morsel.num_rows
 
     if left_rows == 0 or right_rows == 0:
-        return numpy.empty(0, dtype=numpy.int64), numpy.empty(0, dtype=numpy.int64)
+        return (), ()
 
     # Get column vectors
     cdef object left_col_bytes = left_column.encode('utf-8')
@@ -76,28 +68,28 @@ cpdef tuple non_equi_nested_loop_join(
     # Nested loop join - compare each left row with each right row
     for i in range(left_rows):
         left_val = left_vec[i]
-        
+
         # Skip null values in left
         if left_val is None:
             continue
 
         for j in range(right_rows):
             right_val = right_vec[j]
-            
+
             # Skip null values in right
             if right_val is None:
                 continue
 
             # Perform the comparison
-            if comparison_op == 'not_equals':
+            if comparison_op == 'NotEq':
                 comparison_result = left_val != right_val
-            elif comparison_op == 'greater_than':
+            elif comparison_op == 'Gt':
                 comparison_result = left_val > right_val
-            elif comparison_op == 'greater_than_or_equals':
+            elif comparison_op == 'GtEq':
                 comparison_result = left_val >= right_val
-            elif comparison_op == 'less_than':
+            elif comparison_op == 'Lt':
                 comparison_result = left_val < right_val
-            elif comparison_op == 'less_than_or_equals':
+            elif comparison_op == 'LtEq':
                 comparison_result = left_val <= right_val
             else:
                 raise ValueError(f"Unsupported comparison operator: {comparison_op}")

opteryx/compiled/joins/non_equi_join.pyx …piled/joins/nested_loop_join_non_eqi.pyxopteryx/compiled/joins/non_equi_join.pyx renamed to opteryx/compiled/joins/nested_loop_join_non_eqi.pyx opteryx/compiled/joins/non_equi_join.pyx renamed to opteryx/compiled/joins/nested_loop_join_non_eqi.pyx

@@ -32,6 +32,7 @@ cdef class IntBuffer:
     cpdef void append(self, int64_t value)
     cpdef void extend(self, iterable)
     cpdef numpy.ndarray[int64_t, ndim=1] to_numpy(self)
+    cpdef const int64_t[::1] get_buffer(self)
     cpdef size_t size(self)
     cpdef void extend_numpy(self, numpy.ndarray[int64_t, ndim=1] arr)
     cpdef void reserve(self, size_t capacity)

opteryx/compiled/structures/buffers.pxd

@@ -87,6 +87,24 @@ cdef class IntBuffer:
         memcpy(<void*>&arr[0], <const void*>data_ptr, size * sizeof(int64_t))
         return arr
 
+    cpdef const int64_t[::1] get_buffer(self):
+        """
+        Get a read-only memoryview of the underlying buffer (zero-copy).
+
+        This provides direct access to the C++ buffer without copying.
+        The memoryview remains valid as long as the IntBuffer exists
+        and no modifications are made to it.
+
+        Returns:
+            const int64_t[::1]: Read-only memoryview of the buffer
+        """
+        cdef size_t size = self.c_buffer.size()
+        if size == 0:
+            return numpy.empty(0, dtype=numpy.int64)
+
+        cdef const int64_t* data_ptr = self.c_buffer.data()
+        return <const int64_t[:size]>data_ptr
+
     cpdef size_t size(self):
         return self.c_buffer.size()
 

opteryx/compiled/structures/buffers.pyx

@@ -13,6 +13,8 @@
 - evaluate: Compiled expression evaluator
 """
 
+from opteryx.draken.morsels.align import align_tables
+from opteryx.draken.morsels.align import align_tables_pyarray
 from opteryx.draken.morsels.morsel import Morsel
 from opteryx.draken.vectors.vector import Vector
 
@@ -34,4 +36,4 @@ def evaluate(morsel, expression):
     return _evaluate(morsel, expression)
 
 
-__all__ = ("Vector", "Morsel", "evaluate")
+__all__ = ("Vector", "Morsel", "evaluate", "align_tables", "align_tables_pyarray")

opteryx/draken/__init__.py

@@ -3,3 +3,8 @@
 
 Morsel (batch) data structures for columnar processing.
 """
+
+from opteryx.draken.morsels.align import align_tables
+from opteryx.draken.morsels.align import align_tables_pyarray
+
+__all__ = ["align_tables", "align_tables_pyarray"]