feat(python): SAGE integration with Python bindings v0.1.3 by ZeroJustMe · Pull Request #100 · DataSysResearch/BriskFlow

ZeroJustMe · 2026-01-21T14:09:32Z

Summary

This PR adds complete SAGE integration support for SageFlow with enhanced Python bindings.

Changes

New Features

Dual-stream Join Pipeline: Support for Query Stream + Document Stream architecture
SAGE Integration Operators: Python bindings for seamless SAGE DataStream integration
RAG Pipeline Support: Real-time vector similarity join for Retrieval-Augmented Generation

New Files

examples/python/test_sageflow_cpp_runtime.py: Comprehensive C++ runtime verification tests
examples/python/sage_sageflow_dual_stream_join.py: Dual-stream Join pipeline demo
examples/python/llm_inference_service_demo.py: LLM inference service demo
docs/SAGEFLOW_SAGE_INTEGRATION_GUIDE.md: Integration guide

Modified Files

sage_flow/bindings.cpp: Enhanced Python bindings for dual-stream Join
sage_flow/__init__.py: Updated exports
pyproject.toml and sage_flow/_version.py: Version bump to 0.1.3

Testing

All 6 C++ runtime tests pass
Dual-stream Join produces correct similarity matches
RAG context building verified with sample data

Version

0.1.3 (ready for PyPI release)

- Add test_sageflow_cpp_runtime.py: comprehensive C++ runtime verification tests - Add sage_sageflow_dual_stream_join.py: dual-stream Join pipeline demo for RAG - Add SAGEFLOW_SAGE_INTEGRATION_GUIDE.md: integration guide for SAGE + SageFlow The dual-stream Join demo shows: - Query Stream + Document Stream architecture - SageFlow C++ engine for vector similarity join - RAG context building from join results

…ples Modified files: - sage_flow/__init__.py: update exports for SAGE integration - sage_flow/bindings.cpp: enhance Python bindings for dual-stream Join - test/CMakeLists.txt: add new test targets New files: - docs/LLM_INFERENCE_PIPELINE_GUIDE.md: LLM inference pipeline guide - docs/VSJOIN_DESIGN_REVIEW_REPORT.md: VSJoin design review - examples/python/llm_inference_service_demo.py: LLM service demo - examples/python/llm_pipeline_example.py: LLM pipeline example - examples/python/sage_integrated_pipeline_demo.py: SAGE integration demo - test/IntegrationTest/test_non_join_operators_pipeline.cpp: non-join ops test - test/UnitTest/python/: Python unit tests

This release includes: - SAGE integration Python bindings - Dual-stream Join pipeline support - Comprehensive C++ runtime verification tests - RAG pipeline examples and documentation

Copilot

Pull request overview

This PR adds comprehensive SAGE integration support for SageFlow with enhanced Python bindings, enabling dual-stream join pipelines for RAG (Retrieval-Augmented Generation) use cases.

Changes:

Enhanced Python bindings with complete operator API exposure (Join, Filter, Map, Window, Aggregate, TopK, Sink)
Added comprehensive Python unit tests covering API exposure, GIL safety, and pipeline execution
Added C++ integration tests for non-Join operators in multi-threaded execution environment
Version bump to 0.1.3 with new example scripts and integration documentation

Reviewed changes

Copilot reviewed 16 out of 16 changed files in this pull request and generated 21 comments.

Show a summary per file

File	Description
sage_flow/bindings.cpp	Enhanced Python bindings with full operator API, helper functions for numpy conversion, and comprehensive function class bindings
test/UnitTest/python/test_python_bindings.py	New comprehensive Python unit tests for API exposure, vector operations, callbacks, and pipeline execution
test/IntegrationTest/test_non_join_operators_pipeline.cpp	New C++ integration tests validating non-Join operators in ExecutionGraph framework
test/CMakeLists.txt	Added new integration test to build configuration
sage_flow/_version.py	Version bump from 0.1.1.3 to 0.1.3
sage_flow/init.py	Updated exports to include all new function classes and enums
pyproject.toml	Version bump to 0.1.3
examples/python/*.py	New example scripts demonstrating SAGE integration patterns
docs/*.md	New integration guides and design documentation

Copilot · 2026-01-21T14:19:42Z

+
+    pipeline = (
+        query_stream
+        .join(doc_stream, create_combine_vectors_join(), dim=dim, parallelism=1)
+        .writeSink(context_builder.on_match, parallelism=1)


Variable pipeline is not used.

Suggested change

pipeline = (

query_stream

.join(doc_stream, create_combine_vectors_join(), dim=dim, parallelism=1)

.writeSink(context_builder.on_match, parallelism=1)

query_stream.join(

doc_stream,

create_combine_vectors_join(),

dim=dim,

parallelism=1,

).writeSink(

context_builder.on_match,

parallelism=1,

Copilot · 2026-01-21T14:19:43Z

+    aggregator = OnlineAggregator(window_size_ms=3000)
+
+    # 使用 Map 实现在线聚合
+    pipeline = (


Variable pipeline is not used.

Suggested change

pipeline = (

(

Copilot · 2026-01-21T14:19:43Z

+    pipeline = (
+        message_stream
+        .writeSink(on_message, parallelism=1)
+    )


Variable pipeline is not used.

Suggested change

pipeline = (

message_stream

.writeSink(on_message, parallelism=1)

)

message_stream.writeSink(on_message, parallelism=1)

Copilot · 2026-01-21T14:19:43Z

+        print(f"[Sink] Received uid={uid}, ts={ts}, norm={np.linalg.norm(data):.4f}")
+
+    # Chain operators: filter low-norm -> normalize -> collect
+    pipeline = (


Variable pipeline is not used.

Suggested change

pipeline = (

(

Copilot · 2026-01-21T14:19:43Z

+    # Build join pipeline
+    pipeline = (


Variable pipeline is not used.

Suggested change

# Build join pipeline

pipeline = (

# Build join pipeline (side-effect: attach operators to streams)

(

Copilot · 2026-01-21T14:19:46Z

+_SAGE_KERNEL_AVAILABLE = False
+try:
+    from sage.kernel.api import LocalEnvironment
+    from sage.kernel.api.datastream import DataStream


Import of 'DataStream' is not used.

Suggested change

from sage.kernel.api.datastream import DataStream

Copilot · 2026-01-21T14:19:46Z

+    from sage.kernel.api.datastream import DataStream
+    from sage.common.core.functions.map_function import MapFunction
+    from sage.common.core.functions.sink_function import SinkFunction
+    from sage.common.core.functions.source_function import SourceFunction


Import of 'SourceFunction' is not used.

Suggested change

from sage.common.core.functions.source_function import SourceFunction

Copilot · 2026-01-21T14:19:47Z

+_SAGE_EMBEDDING_AVAILABLE = False
+try:
+    from sage.common.components.sage_embedding import (
+        EmbeddingClientAdapter,


Import of 'EmbeddingClientAdapter' is not used.

Suggested change

EmbeddingClientAdapter,

Copilot · 2026-01-21T14:19:47Z

+        SageFlowAggregationOperator,
+    )
+    _SAGE_FLOW_OPERATORS_AVAILABLE = True
+    print("[Setup] ✓ SAGE SageFlow Operators (sage-middleware) 可用")


Import of 'SageFlowJoinOperator' is not used.
Import of 'SageFlowAggregationOperator' is not used.

Suggested change

print("[Setup] ✓ SAGE SageFlow Operators (sage-middleware) 可用")

print("[Setup] ✓ SAGE SageFlow Operators (sage-middleware) 可用")

# Mark imports as intentionally used so that static analyzers do not flag them.

_SAGE_FLOW_IMPORTED_OPERATORS = (SageFlowJoinOperator, SageFlowAggregationOperator)

Copilot · 2026-01-21T14:19:47Z

+from sage.common.core.functions.map_function import MapFunction
+from sage.common.core.functions.sink_function import SinkFunction
+from sage.common.core.functions.comap_function import BaseCoMapFunction
+from sage.kernel.api.local_environment import LocalEnvironment


Import of 'LocalEnvironment' is not used.

Suggested change

from sage.kernel.api.local_environment import LocalEnvironment

The inferDefaults() for LSH algorithm returns PARTITIONED, not PARTITIONED_VECTOR. PARTITIONED_VECTOR is only used for VSJOIN algorithm.

LSH algorithm doesn't use external index, so left_index_id and right_index_id can be -1. Removed these checks to match feat/implement_vsjoin branch.

* feat(python): add SAGE integration examples and documentation - Add test_sageflow_cpp_runtime.py: comprehensive C++ runtime verification tests - Add sage_sageflow_dual_stream_join.py: dual-stream Join pipeline demo for RAG - Add SAGEFLOW_SAGE_INTEGRATION_GUIDE.md: integration guide for SAGE + SageFlow The dual-stream Join demo shows: - Query Stream + Document Stream architecture - SageFlow C++ engine for vector similarity join - RAG context building from join results * feat(python): complete SAGE integration with Python bindings and examples Modified files: - sage_flow/__init__.py: update exports for SAGE integration - sage_flow/bindings.cpp: enhance Python bindings for dual-stream Join - test/CMakeLists.txt: add new test targets New files: - docs/LLM_INFERENCE_PIPELINE_GUIDE.md: LLM inference pipeline guide - docs/VSJOIN_DESIGN_REVIEW_REPORT.md: VSJoin design review - examples/python/llm_inference_service_demo.py: LLM service demo - examples/python/llm_pipeline_example.py: LLM pipeline example - examples/python/sage_integrated_pipeline_demo.py: SAGE integration demo - test/IntegrationTest/test_non_join_operators_pipeline.cpp: non-join ops test - test/UnitTest/python/: Python unit tests * chore: bump version to 0.1.3 for PyPI release This release includes: - SAGE integration Python bindings - Dual-stream Join pipeline support - Comprehensive C++ runtime verification tests - RAG pipeline examples and documentation * fix(test): correct LSH window_state_type expectation to PARTITIONED The inferDefaults() for LSH algorithm returns PARTITIONED, not PARTITIONED_VECTOR. PARTITIONED_VECTOR is only used for VSJOIN algorithm. * fix(test): remove index_id checks from LSH test LSH algorithm doesn't use external index, so left_index_id and right_index_id can be -1. Removed these checks to match feat/implement_vsjoin branch. * doc: add sage pipeline markdown file

* feat(vsjoin): Refine VSJoin plan and add detailed task docs This commit introduces a comprehensive plan and task breakdown for the new VSJoin implementation. - **Refine VSJoin Plan:** - The main design document is updated to unify naming (removing "v2"). - Clarifies the replacement of v1 components with the new architecture (TwoTierWindowState + ConcurrencyManager). - Adds detailed sections on UID deduplication and the RCU-based load balancing mechanism (AssignmentTable). - **Add Detailed Task Documents:** - Creates a new directory. - Adds markdown files for each implementation step (Task 01 to 09), providing clear instructions for development. These files are force-added as the parent directory is in .gitignore. - **Cleanup Unused v1 Components:** - Deletes and as they are no longer needed in the new design. * feat(test): Add unit test for VSJoin rebuild functionality This commit adds a new unit test for the VSJoin rebuild process, enhancing the test coverage for the VSJoin implementation. The test file is located at UnitTest/test_vsjoin_rebuild.cpp and is configured to run with a specified timeout of 300 seconds. * feat(vsjoin): implement VSJoin dual-index architecture (tasks 01-04) - Task 01: VSJoinMethod basic implementation - Implement ExecuteEager() with dual-layer query logic (Global + Local) - Add UID deduplication using local unordered_set - Implement setGlobalIndexIds/setLocalIndexIds/setWindowStates interfaces - Task 02: JoinStrategyFactory integration - Add JoinAlgorithm::VSJOIN enum - Add vsjoin_* configuration parameters - Integrate VSJOIN case in factory (fallback to BruteForce for now) - Task 03: JoinOperator VSJoin special path - Add vsjoin_local_*_ids_ and vsjoin_global_*_id_ members - Implement VSJoin-specific updateSideWithState() (insert to local index only) - Support LSH partitioner for VSJoin - Task 04: Background rebuild mechanism - Implement globalIndexRebuildLoop() with periodic rebuild - Use std::call_once for thread-safe single startup - Implement local unordered_set deduplication (lock-free) - Add atomic index replacement via replace_index_by_id() Integration tests passed: - bruteforce: 7/7 tests, recall=1.000 - ivf: 7/7 tests, recall=0.999 - hdr_tree: 7/7 tests, recall=1.000 - clustered_join: 3/3 tests, recall=1.000 * fix(vsjoin): fix dangling pointer bug in globalIndexRebuildLoop The snapshot ownership was released prematurely in the rebuild loop, causing the pointers stored in unique_left_records and unique_right_records to become dangling. This led to IVF centroids being initialized with garbage data (dimension=0), causing 'Vectors must be of the same size' errors during query_for_join. Fix: Keep snapshot vectors alive until build_index_from_records completes by storing them in left_snapshots/right_snapshots containers. Also includes: - Improved test assertions for query record dimension - Code cleanup for VSJoin factory integration * chore: bump version to 0.1.3.1 * chore: add TODO issue links via todo-to-issue-action * fix(test): correct LSH window_state_type expectation to PARTITIONED The inferDefaults() for LSH algorithm returns PARTITIONED, not PARTITIONED_VECTOR. PARTITIONED_VECTOR is only used for VSJOIN algorithm. * feat(vsjoin): 完成 Task 06 - VSJoin 集成测试链路打通 sed -n '168,195p' /root/sageFlow/src/operator/utils/join_strategy_factory.cpp 1. 配置验证 ✅ - integration_test_cases.toml 包含 4 个启用的 VSJoin 测试用例 - 配置包含必要参数 (vsjoin_num_hash_functions, vsjoin_boundary_threshold 等) - num_partitions 参数设置合理 2. 链路打通验证 ✅ - JoinStrategyFactory::create() 正确创建 VSJoinMethod - TwoTierWindowState 正确初始化 - Global/Local Index 正确创建和管理 - 后台重建线程正常工作 3. 测试执行验证 ✅ - test_join_baseline_integration --gtest_filter='*vsjoin*' 执行成功 - run_integration_test.py --methods vsjoin 执行成功 - 测试报告正确生成 4. 召回率验证 ✅ - vsjoin_baseline: Recall=1.0 (预期>=0.70) - vsjoin_high_recall: Recall=1.0 (预期>=0.75) - vsjoin_parallelism_scaling: Recall=1.0 (并行度 1-16) - vsjoin_low_latency: Recall>=0.60 sed -n '168,195p' /root/sageFlow/src/operator/utils/join_strategy_factory.cpp - 修复 JoinConfigValidator 允许 LSH + TWO_TIER 组合 - 修复 JoinOperator 中 VSJoin 的 use_index_ 和 index_id 计算 - 更新 integration_test_cases.toml 添加 VSJoin 测试用例 - 更新 task06_integration_test.md 添加集成测试框架链路任务 * feat(vsjoin): implement Task 07-09 - AssignmentTable, LoadMonitor, Logical Partition Routing, and Load Balancing Tests Task 07: AssignmentTable (RCU) + LoadMonitor - Implement PartitionAssignment with RCU pattern for lock-free reads - Implement LoadMonitor for tracking partition load statistics - Support logical-to-physical partition mapping - Batch atomic updates for assignment table Task 08: Logical Partition Routing Integration - Add VSJoin routing methods in JoinOperator - Implement routeByLSHBucket() for query routing - Implement determineTargetPartitions() with multi-partition support - Integrate with CentroidPartitioner for initial assignment Task 09: Load Balancing Tests - Add comprehensive unit tests for LoadMonitor - Add comprehensive unit tests for PartitionAssignment - Add integration tests for VSJoin routing - Add load balancing scenario tests * test: isolate per-run outputs (reports/logs/charts) * feat(python): SAGE integration with Python bindings v0.1.3 (#100) * feat(python): add SAGE integration examples and documentation - Add test_sageflow_cpp_runtime.py: comprehensive C++ runtime verification tests - Add sage_sageflow_dual_stream_join.py: dual-stream Join pipeline demo for RAG - Add SAGEFLOW_SAGE_INTEGRATION_GUIDE.md: integration guide for SAGE + SageFlow The dual-stream Join demo shows: - Query Stream + Document Stream architecture - SageFlow C++ engine for vector similarity join - RAG context building from join results * feat(python): complete SAGE integration with Python bindings and examples Modified files: - sage_flow/__init__.py: update exports for SAGE integration - sage_flow/bindings.cpp: enhance Python bindings for dual-stream Join - test/CMakeLists.txt: add new test targets New files: - docs/LLM_INFERENCE_PIPELINE_GUIDE.md: LLM inference pipeline guide - docs/VSJOIN_DESIGN_REVIEW_REPORT.md: VSJoin design review - examples/python/llm_inference_service_demo.py: LLM service demo - examples/python/llm_pipeline_example.py: LLM pipeline example - examples/python/sage_integrated_pipeline_demo.py: SAGE integration demo - test/IntegrationTest/test_non_join_operators_pipeline.cpp: non-join ops test - test/UnitTest/python/: Python unit tests * chore: bump version to 0.1.3 for PyPI release This release includes: - SAGE integration Python bindings - Dual-stream Join pipeline support - Comprehensive C++ runtime verification tests - RAG pipeline examples and documentation * fix(test): correct LSH window_state_type expectation to PARTITIONED The inferDefaults() for LSH algorithm returns PARTITIONED, not PARTITIONED_VECTOR. PARTITIONED_VECTOR is only used for VSJOIN algorithm. * fix(test): remove index_id checks from LSH test LSH algorithm doesn't use external index, so left_index_id and right_index_id can be -1. Removed these checks to match feat/implement_vsjoin branch. * doc: add sage pipeline markdown file * fix: unitest failed * fix: update VSJoin test to expect TWO_TIER window state type * fix(pybind): fix alpha hardcode bug in NORMALIZED mode and add window_size_ms parameter - Fix BruteForceBaseline to use configured similarity_alpha instead of hardcoded 0.1 in NORMALIZED mode - Add 7-parameter join() overload with window_size_ms support in pybind interface - Add StreamingSource for dynamic streaming input (vs batch-oriented SimpleStreamSource) - Update JoinStrategyConfig to properly propagate similarity_alpha and window_size_ms * refactor(pybind): simplify join method bindings and fix test expectations - Remove unused include of join_strategy_config.h in bindings.cpp - Simplify join binding: use JoinFunction constructor without time_window param - Remove redundant SimilarityMode::NORMALIZED and alpha settings - Fix test expectation: LSH uses PARTITIONED instead of PARTITIONED_VECTOR - Add note about LSH not depending on external index * add refactor partitioner doc, stash changes --------- Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

ZeroJustMe added 3 commits January 19, 2026 12:05

chore: bump version to 0.1.3 for PyPI release

1fcd7ee

This release includes: - SAGE integration Python bindings - Dual-stream Join pipeline support - Comprehensive C++ runtime verification tests - RAG pipeline examples and documentation

Copilot AI review requested due to automatic review settings January 21, 2026 14:09

Copilot started reviewing on behalf of ZeroJustMe January 21, 2026 14:10 View session

Copilot AI reviewed Jan 21, 2026

View reviewed changes

ZeroJustMe added 3 commits January 21, 2026 14:39

fix(test): correct LSH window_state_type expectation to PARTITIONED

b28582b

The inferDefaults() for LSH algorithm returns PARTITIONED, not PARTITIONED_VECTOR. PARTITIONED_VECTOR is only used for VSJOIN algorithm.

fix(test): remove index_id checks from LSH test

620d457

LSH algorithm doesn't use external index, so left_index_id and right_index_id can be -1. Removed these checks to match feat/implement_vsjoin branch.

doc: add sage pipeline markdown file

f85d6f3

ZeroJustMe merged commit 358e387 into main-dev Jan 27, 2026
3 checks passed

Navigation Menu

Uh oh!

feat(python): SAGE integration with Python bindings v0.1.3 - #100

feat(python): SAGE integration with Python bindings v0.1.3#100
ZeroJustMe merged 6 commits into
main-devfrom
feat/pybinding_api_for_sage_service

ZeroJustMe commented Jan 21, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Jan 21, 2026

Uh oh!

Copilot AI Jan 21, 2026

Uh oh!

Copilot AI Jan 21, 2026

Uh oh!

Copilot AI Jan 21, 2026

Uh oh!

Copilot AI Jan 21, 2026

Uh oh!

Copilot AI Jan 21, 2026

Uh oh!

Copilot AI Jan 21, 2026

Uh oh!

Copilot AI Jan 21, 2026

Uh oh!

Copilot AI Jan 21, 2026

Uh oh!

Copilot AI Jan 21, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

ZeroJustMe commented Jan 21, 2026

Summary

Changes

New Features

New Files

Modified Files

Testing

Version

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants