Paging iterator optimization: +1 lookahead to eliminate terminal out-of-range queries, env validation, setuptools bump

Author: FanaticPythoner

.gitignore

@@ -226,3 +226,4 @@ atp_pipeline-1.0.0-cp312-cp312-manylinux_2_34_x86_64.whl
 atp_pipeline-1.0.6-cp312-cp312-manylinux_2_35_x86_64.whl
 setup_env.sh.old
 atp_pipeline-1.0.7-cp312-cp312-manylinux_2_34_x86_64.whl
+TODOS_LISTS/AUDIT-KUZU-SEGFAULT-PAGING.md

pyproject.toml

@@ -1,10 +1,10 @@
 [build-system]
-requires = ["setuptools>=61.0", "wheel"]
+requires = ["setuptools>=82.0", "wheel"]
 build-backend = "setuptools.build_meta"
 
 [project]
 name = "kuzualchemy"
-version = "0.3.22"
+version = "0.3.23"
 description = "SQLAlchemy-like ORM for Kuzu graph database"
 readme = "README.md"
 license = { file = "LICENSE" }

src/kuzualchemy/kuzu_query.py

@@ -31,6 +31,29 @@
 
 logger = logging.getLogger(__name__)
 
+_ENV_ATP_READONLY_POOL_MAX_SIZE = "ATP_READONLY_POOL_MAX_SIZE"
+
+
+def _read_required_positive_int_env(var_name: str) -> int:
+    """Read and validate a required positive integer environment variable."""
+    raw = os.getenv(var_name)
+    if raw is None:
+        raise RuntimeError(
+            f"Missing required environment variable '{var_name}'. "
+            "Configure it before calling Query.iter()."
+        )
+    try:
+        value = int(raw)
+    except ValueError as exc:
+        raise ValueError(
+            f"Environment variable '{var_name}' must be an integer, got: {raw!r}"
+        ) from exc
+    if value <= 0:
+        raise ValueError(
+            f"Environment variable '{var_name}' must be > 0, got: {value}"
+        )
+    return value
+
 if TYPE_CHECKING:
     from .kuzu_session import KuzuSession
 
@@ -599,8 +622,8 @@ def iter(self, page_size: int = 10, prefetch_pages: int = 1) -> Iterator[Union[M
         )
 
         # Check if parallel execution is available and beneficial
-        pool_size = int(os.environ["ATP_READONLY_POOL_MAX_SIZE"])
-        parallel_threshold = int(os.environ["ATP_READONLY_POOL_MAX_SIZE"])
+        pool_size = _read_required_positive_int_env(_ENV_ATP_READONLY_POOL_MAX_SIZE)
+        parallel_threshold = pool_size
         use_parallel = pool_size > 1 and parallel_threshold > 0
 
         def fetch_page(offset: int) -> List[Union[ModelType, Dict[str, Any]]]:
@@ -629,6 +652,37 @@ def fetch_page(offset: int) -> List[Union[ModelType, Dict[str, Any]]]:
 
             return mapped
 
+        def fetch_page_with_lookahead(offset: int) -> Tuple[List[Union[ModelType, Dict[str, Any]]], bool]:
+            """Fetch a page with one-row lookahead to avoid terminal out-of-range SKIP queries."""
+            q = self.offset(offset).limit(ps + 1)
+            t0 = time.perf_counter()
+            raw = q._execute()
+            t1 = time.perf_counter()
+            mapped = q._map_results(raw)
+            t2 = time.perf_counter()
+
+            has_more = len(mapped) > ps
+            page_data = mapped[:ps] if has_more else mapped
+
+            if getattr(self._session, "_debug_timing", False) or ((t2 - t0) >= 0.25):
+                raw_rows = len(raw) if isinstance(raw, list) else None
+                mapped_rows = len(page_data) if isinstance(page_data, list) else None
+                logger.info(
+                    "kuzu.query.page.lookahead rel=%s offset=%d page_size=%d raw_rows=%s mapped_rows=%s has_more=%s exec_seconds=%.6f map_seconds=%.6f total_seconds=%.6f pairs_subset=%s",
+                    model_name,
+                    int(offset),
+                    int(ps),
+                    raw_rows,
+                    mapped_rows,
+                    has_more,
+                    (t1 - t0),
+                    (t2 - t1),
+                    (t2 - t0),
+                    pairs_subset_meta,
+                )
+
+            return page_data, has_more
+
         def fetch_pages_parallel(offsets: List[int]) -> List[List[Union[ModelType, Dict[str, Any]]]]:
             """Fetch multiple pages in parallel using Rust rayon via ATP pipeline."""
             if not offsets:
@@ -672,63 +726,81 @@ def fetch_pages_parallel(offsets: List[int]) -> List[List[Union[ModelType, Dict[
                 mapped_pages.append(mapped)
             return mapped_pages
 
-        # Fetch first page to determine if more pages exist
-        offset = 0
-        page = fetch_page(offset)
-        offset += ps
+        # If parallel execution is enabled, preserve existing count-bounded parallel strategy.
+        if use_parallel:
+            offset = 0
+            page = fetch_page(offset)
+            offset += ps
+
+            # If first page is not full, result set fits in one page.
+            if len(page) < ps:
+                for item in page:
+                    yield item
+                return
+
+            total_rows = self.count_results()
+            remaining_rows = max(total_rows - ps, 0)
 
-        # If parallel execution is enabled and first page is full, try parallel fetching
-        if use_parallel and len(page) == ps:
             # Yield first page items
             for item in page:
                 yield item
-            
+
+            if remaining_rows == 0:
+                return
+
             # Parallel batch fetching
             batch_size = min(pool_size, parallel_threshold)
-            while True:
-                # Build batch of offsets
-                batch_offsets = [offset + i * ps for i in range(batch_size)]
-                
+            while remaining_rows > 0:
+                pages_in_batch = min(batch_size, (remaining_rows + ps - 1) // ps)
+                batch_offsets = [offset + i * ps for i in range(pages_in_batch)]
+
                 # Fetch batch in parallel
                 batch_pages = fetch_pages_parallel(batch_offsets)
-                
-                # Yield results and track if we got a partial page
-                last_page_full = True
-                for page_idx, page_data in enumerate(batch_pages):
+
+                # Yield results in requested page order
+                for page_data in batch_pages:
                     for item in page_data:
                         yield item
-                    if len(page_data) < ps:
-                        last_page_full = False
-                        break
-                
-                if not last_page_full:
-                    break
-                
-                offset += batch_size * ps
-        elif pf > 0:
+
+                advanced_rows = pages_in_batch * ps
+                offset += advanced_rows
+                remaining_rows = max(remaining_rows - advanced_rows, 0)
+
+            return
+
+        # Sequential modes: use +1 lookahead to avoid issuing a terminal out-of-range page.
+        offset = 0
+        page, has_more = fetch_page_with_lookahead(offset)
+        offset += ps
+
+        if pf > 0:
             # Sequential with prefetch (original behavior)
             with ThreadPoolExecutor(max_workers=1) as executor:
-                next_future = executor.submit(fetch_page, offset) if len(page) == ps else None
+                next_future = executor.submit(fetch_page_with_lookahead, offset) if has_more else None
                 while True:
                     for item in page:
                         yield item
-                    if len(page) < ps:
+                    if not has_more:
                         break
-                    next_page = next_future.result() if next_future is not None else fetch_page(offset)
+                    if next_future is not None:
+                        next_page, next_has_more = next_future.result()
+                    else:
+                        next_page, next_has_more = fetch_page_with_lookahead(offset)
                     offset += ps
-                    if len(next_page) == ps and pf > 0:
-                        next_future = executor.submit(fetch_page, offset)
+                    if next_has_more and pf > 0:
+                        next_future = executor.submit(fetch_page_with_lookahead, offset)
                     else:
                         next_future = None
                     page = next_page
+                    has_more = next_has_more
         else:
             # Pure sequential (no prefetch)
             while True:
                 for item in page:
                     yield item
-                if len(page) < ps:
+                if not has_more:
                     break
-                page = fetch_page(offset)
+                page, has_more = fetch_page_with_lookahead(offset)
                 offset += ps
 
     def all(self, as_iterator: bool = False, page_size: Optional[int] = None, prefetch_pages: int = 1) -> Union[List[ModelType], List[Dict[str, Any]], Iterator[Union[ModelType, Dict[str, Any]]]]:
@@ -784,7 +856,9 @@ def exists(self) -> bool:
 
     def count_results(self) -> int:
         """Count the number of results."""
-        count_query = self.count()
+        # ORDER BY columns are not valid after scalar COUNT aggregation in Kuzu.
+        # Keep all filters/joins while stripping ORDER BY for the COUNT query only.
+        count_query = self._copy_with_state(order_by=[]).count()
         result = count_query._execute()
         if type(result) is not list:
             logger.error("Count query returned non-list result type: %r", type(result))

src/kuzualchemy/kuzu_session.py

@@ -265,37 +265,44 @@ def iterate(
         if " skip " in f" {q_lower} " or " limit " in f" {q_lower} ":
             raise ValueError("Raw query already contains SKIP/LIMIT; cannot auto-paginate. Remove them and retry.")
 
-        def fetch_page(offset: int) -> List[Dict[str, Any]]:
-            paged_q = f"{base} SKIP {offset} LIMIT {page_size}"
-            return self.execute(paged_q, parameters)
+        def fetch_page(offset: int) -> tuple[List[Dict[str, Any]], bool]:
+            paged_q = f"{base} SKIP {offset} LIMIT {page_size + 1}"
+            rows = self.execute(paged_q, parameters)
+            has_more = len(rows) > page_size
+            page = rows[:page_size] if has_more else rows
+            return page, has_more
 
         # First page
         offset = 0
-        page = fetch_page(offset)
+        page, has_more = fetch_page(offset)
         if prefetch_pages > 0:
             from concurrent.futures import ThreadPoolExecutor
             with ThreadPoolExecutor(max_workers=1) as executor:
-                next_future = executor.submit(fetch_page, offset + page_size) if len(page) == page_size else None
+                next_future = executor.submit(fetch_page, offset + page_size) if has_more else None
                 while True:
                     for row in page:
                         yield row
-                    if len(page) < page_size:
+                    if not has_more:
                         break
-                    next_page = next_future.result() if next_future is not None else fetch_page(offset + page_size)
+                    if next_future is not None:
+                        next_page, next_has_more = next_future.result()
+                    else:
+                        next_page, next_has_more = fetch_page(offset + page_size)
                     offset += page_size
-                    if len(next_page) == page_size:
+                    if next_has_more:
                         next_future = executor.submit(fetch_page, offset + page_size)
                     else:
                         next_future = None
                     page = next_page
+                    has_more = next_has_more
         else:
             while True:
                 for row in page:
                     yield row
-                if len(page) < page_size:
+                if not has_more:
                     break
                 offset += page_size
-                page = fetch_page(offset)
+                page, has_more = fetch_page(offset)
 
 
     def _execute_with_connection_reuse(self, query: str, parameters: Optional[Dict[str, Any]] = None) -> Any:

tests/.env.test

@@ -0,0 +1,11 @@
+# Kuzualchemy pytest-specific environment.
+# Keys are namespaced to avoid cross-project collisions when other repositories
+# define similarly named runtime variables.
+
+KUZUALCHEMY_TEST_OMP_NUM_THREADS=1
+KUZUALCHEMY_TEST_MKL_NUM_THREADS=1
+KUZUALCHEMY_TEST_NUMBA_NUM_THREADS=1
+KUZUALCHEMY_TEST_ATP_PROFILE=0
+KUZUALCHEMY_TEST_ATP_READONLY_POOL_MAX_SIZE=1
+KUZUALCHEMY_TEST_ATP_READONLY_POOL_WARM_COUNT=1
+KUZUALCHEMY_TEST_ATP_READONLY_BUFFER_POOL_BYTES=0

tests/_env.py

@@ -1,23 +1,76 @@
 """Test environment loader for Kuzualchemy pytest suite.
 
-This module ensures that the repository-level `.env` file is loaded before any
-tests run. All parsing and environment application is implemented in Rust in
-the `atp_core::env` module; here we simply delegate to the public
-``atp_pipeline.load_workspace_dotenv`` helper so Rust-only and Python-driven
-entrypoints share identical semantics.
+This module loads a dedicated test environment file (`tests/.env.test`) before
+any tests run. Keys in the file are intentionally namespaced to avoid
+cross-project collisions when other repositories expose similarly named runtime
+variables.
 """
 
 from __future__ import annotations
 
 import os
+from pathlib import Path
 
-from atp_pipeline import load_workspace_dotenv
 
+_TEST_ENV_PATH = Path(__file__).with_name(".env.test")
 
-load_workspace_dotenv(required=True)
+_TEST_TO_RUNTIME_ENV_KEY_MAP = {
+    "KUZUALCHEMY_TEST_OMP_NUM_THREADS": "OMP_NUM_THREADS",
+    "KUZUALCHEMY_TEST_MKL_NUM_THREADS": "MKL_NUM_THREADS",
+    "KUZUALCHEMY_TEST_NUMBA_NUM_THREADS": "NUMBA_NUM_THREADS",
+    "KUZUALCHEMY_TEST_ATP_PROFILE": "ATP_PROFILE",
+    "KUZUALCHEMY_TEST_ATP_READONLY_POOL_MAX_SIZE": "ATP_READONLY_POOL_MAX_SIZE",
+    "KUZUALCHEMY_TEST_ATP_READONLY_POOL_WARM_COUNT": "ATP_READONLY_POOL_WARM_COUNT",
+    "KUZUALCHEMY_TEST_ATP_READONLY_BUFFER_POOL_BYTES": "ATP_READONLY_BUFFER_POOL_BYTES",
+}
 
-os.environ["ATP_PROFILE"] = "0"
-os.environ.pop("ATP_PROFILE_FREQ", None)
-os.environ.pop("ATP_PROFILE_FLAMEGRAPH", None)
-os.environ.pop("ATP_PROFILE_SPEEDSCOPE", None)
+
+def _parse_env_file(path: Path) -> dict[str, str]:
+    if not path.exists():
+        raise RuntimeError(f"Missing required test env file: {path}")
+
+    parsed: dict[str, str] = {}
+    lines = path.read_text(encoding="utf-8").splitlines()
+    for line_no, raw_line in enumerate(lines, start=1):
+        line = raw_line.strip()
+        if not line or line.startswith("#"):
+            continue
+        if "=" not in line:
+            raise ValueError(
+                f"Invalid env entry in {path} at line {line_no}: expected KEY=VALUE"
+            )
+        key_part, value_part = line.split("=", 1)
+        key = key_part.strip()
+        value = value_part.strip()
+        if not key:
+            raise ValueError(
+                f"Invalid env entry in {path} at line {line_no}: empty KEY"
+            )
+        if key in parsed:
+            raise ValueError(
+                f"Duplicate env key in {path} at line {line_no}: {key}"
+            )
+        parsed[key] = value
+    return parsed
+
+
+def _apply_runtime_test_env(parsed_env: dict[str, str]) -> None:
+    missing_required_keys = [
+        key for key in _TEST_TO_RUNTIME_ENV_KEY_MAP if key not in parsed_env
+    ]
+    if missing_required_keys:
+        missing_csv = ", ".join(sorted(missing_required_keys))
+        raise RuntimeError(
+            f"Missing required keys in {_TEST_ENV_PATH}: {missing_csv}"
+        )
+
+    for source_key, target_key in _TEST_TO_RUNTIME_ENV_KEY_MAP.items():
+        os.environ[target_key] = parsed_env[source_key]
+
+    os.environ.pop("ATP_PROFILE_FREQ", None)
+    os.environ.pop("ATP_PROFILE_FLAMEGRAPH", None)
+    os.environ.pop("ATP_PROFILE_SPEEDSCOPE", None)
+
+
+_apply_runtime_test_env(_parse_env_file(_TEST_ENV_PATH))