Streaming Q4 implementation #710

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open

TomAugspurger wants to merge 85 commits into rapidsai:main from TomAugspurger:tom/streaming-q4

cpp/benchmarks/streaming/ndsh/CMakeLists.txt

-Original file line number
+Diff line change
@@ Expand Up / @@ -41,7 +41,7 @@ target_link_libraries( @@
                             $<TARGET_NAME_IF_EXISTS:conda_env> maybe_asan
     )
-    set(RAPIDSMPFNDSH_QUERIES q01 q03 q09 q21 bench_read)
+    set(RAPIDSMPFNDSH_QUERIES q01 q03 q04 q09 q21 bench_read)
     foreach(query IN ITEMS ${RAPIDSMPFNDSH_QUERIES})
       add_executable(${query} "${query}.cpp")
@@ Expand Down @@

cpp/benchmarks/streaming/ndsh/join.cpp

-Original file line number
+Diff line change
@@ Expand Up / @@ -13,6 +13,7 @@ @@
     #include <cudf/concatenate.hpp>
     #include <cudf/contiguous_split.hpp>
     #include <cudf/copying.hpp>
+    #include <cudf/join/filtered_join.hpp>
     #include <cudf/join/hash_join.hpp>
     #include <cudf/table/table.hpp>
     #include <cudf/table/table_view.hpp>
@@ Expand Down Expand Up / @@ -144,6 +145,75 @@ streaming::Node broadcast( @@
         co_await ch_out->drain(ctx->executor());
     }
+    /**
+     * @brief Join a table chunk against a build hash table returning a message of the result.
+     *
+     * @param ctx Streaming context
+     * @param left_chunk Chunk to join. Used as the probe table in a filtered join.
+     * @param right_chunk Chunk to join. Used as the build table in a filtered join.
+     * @param left_carrier Columns from `left_chunk` to include in the output.
+     * @param left_on Key column indices in `left_chunk`.
+     * @param right_on Key column indices in `right_chunk`.
+     * @param sequence Sequence number of the output
+     * @param left_event Event recording the availability of `left_chunk`.
+     *
+     * @return Message of `TableChunk` containing the result of the semi join.
+     */
+    streaming::Message semi_join_chunk(
+        std::shared_ptr<streaming::Context> ctx,
+        streaming::TableChunk const& left_chunk,
+        streaming::TableChunk&& right_chunk,
+        cudf::table_view left_carrier,
+        std::vector<cudf::size_type> left_on,
+        std::vector<cudf::size_type> right_on,
+        std::uint64_t sequence,
+        CudaEvent* left_event
+    ) {
+        auto chunk_stream = right_chunk.stream();
+        left_event->stream_wait(chunk_stream);
+        // At this point, both left_chunk and right_chunk are valid on
+        // either stream. We'll do everything from here out on the
+        // right_chunk.stream(), so that we don't introduce false dependencies
+        // between the different chunks.
+        auto joiner = cudf::filtered_join(
+            right_chunk.table_view().select(right_on),
+            cudf::null_equality::UNEQUAL,
+            cudf::set_as_build_table::RIGHT,
+            chunk_stream
+        );
+        auto match = joiner.semi_join(
+            left_chunk.table_view().select(left_on), chunk_stream, ctx->br()->device_mr()
+        );
+        ctx->comm()->logger().debug(
+            "semi_join_chunk: left.num_rows()=", left_chunk.table_view().num_rows()
+        );
+        ctx->comm()->logger().debug("semi_join_chunk: match.size()=", match->size());
+        cudf::column_view indices = cudf::device_span<cudf::size_type const>(*match);
+        auto result_columns = cudf::gather(
+                                  left_carrier,
+                                  indices,
+                                  cudf::out_of_bounds_policy::DONT_CHECK,
+                                  chunk_stream,
+                                  ctx->br()->device_mr()
+        )
+                                  ->release();
+        auto result_table = std::make_unique<cudf::table>(std::move(result_columns));
+        // Deallocation of the join indices will happen on chunk_stream, so add stream dep
+        cuda_stream_join(left_chunk.stream(), chunk_stream);
+        return streaming::to_message(
+            sequence,
+            std::make_unique<streaming::TableChunk>(std::move(result_table), chunk_stream)
+        );
+    }
     /**
      * @brief Join a table chunk against a build hash table returning a message of the result.
      *
@@ Expand Down Expand Up / @@ -345,6 +415,131 @@ streaming::Node inner_join_shuffle( @@
         co_await ch_out->drain(ctx->executor());
     }
+    streaming::Node left_semi_join_broadcast_left(
+        std::shared_ptr<streaming::Context> ctx,
+        std::shared_ptr<streaming::Channel> left,
+        std::shared_ptr<streaming::Channel> right,
+        std::shared_ptr<streaming::Channel> ch_out,
+        std::vector<cudf::size_type> left_on,
+        std::vector<cudf::size_type> right_on,
+        OpID tag,
+        KeepKeys keep_keys
+    ) {
+        streaming::ShutdownAtExit c{left, right, ch_out};
+        co_await ctx->executor()->schedule();
+        ctx->comm()->logger().print("Left semi broadcast join ", static_cast<int>(tag));
+        auto left_table = co_await (co_await broadcast(ctx, left, tag))
+                              .release<streaming::TableChunk>()
+                              .make_available(ctx);
+        ctx->comm()->logger().print(
+            "Left (probe) table has ", left_table.table_view().num_rows(), " rows"
+        );
+        CudaEvent left_event;
+        left_event.record(left_table.stream());
+        cudf::table_view left_carrier;
+        if (keep_keys == KeepKeys::YES) {
+            left_carrier = left_table.table_view();
+        } else {
+            std::vector<cudf::size_type> to_keep;
+            std::ranges::copy_if(
+                std::ranges::iota_view(0, left_table.table_view().num_columns()),
+                std::back_inserter(to_keep),
+                [&](auto i) { return std::ranges::find(left_on, i) == left_on.end(); }
+            );
+            left_carrier = left_table.table_view().select(to_keep);
+        }
+        while (!ch_out->is_shutdown()) {
+            auto right_msg = co_await right->receive();
+            if (right_msg.empty()) {
+                break;
+            }
+            // The ``right`` table has been hash-partitioned (via a shuffle) on
+            // the join key. Thanks to the hash-partitioning, we don't need to worry
+            // about deduplicating matches across partitions. Anything that matches
+            // in the semi-join belongs in the output.
+            auto right_chunk =
+                co_await right_msg.release<streaming::TableChunk>().make_available(ctx);
+            co_await ch_out->send(semi_join_chunk(
+                ctx,
+                left_table,
+                std::move(right_chunk),
+                left_carrier,
+                left_on,
+                right_on,
+                right_msg.sequence_number(),
+                &left_event
+            ));
+        }
+        co_await ch_out->drain(ctx->executor());
+    }
+    streaming::Node left_semi_join_shuffle(
+        std::shared_ptr<streaming::Context> ctx,
+        std::shared_ptr<streaming::Channel> left,
+        std::shared_ptr<streaming::Channel> right,
+        std::shared_ptr<streaming::Channel> ch_out,
+        std::vector<cudf::size_type> left_on,
+        std::vector<cudf::size_type> right_on,
+        KeepKeys keep_keys
+    ) {
+        streaming::ShutdownAtExit c{left, right, ch_out};
+        ctx->comm()->logger().print("Shuffle left semi join");
+        co_await ctx->executor()->schedule();
+        CudaEvent left_event;
+        while (!ch_out->is_shutdown()) {
+            // Requirement: two shuffles kick out partitions in the same order
+            auto left_msg = co_await left->receive();
+            auto right_msg = co_await right->receive();
+            if (left_msg.empty()) {
+                RAPIDSMPF_EXPECTS(
+                    right_msg.empty(), "Left does not have same number of partitions as right"
+                );
+                break;
+            }
+            RAPIDSMPF_EXPECTS(
+                left_msg.sequence_number() == right_msg.sequence_number(),
+                "Mismatching sequence numbers"
+            );
+            auto left_chunk =
+                co_await left_msg.release<streaming::TableChunk>().make_available(ctx);
+            auto right_chunk =
+                co_await right_msg.release<streaming::TableChunk>().make_available(ctx);
+            left_event.record(left_chunk.stream());
+            cudf::table_view left_carrier;
+            if (keep_keys == KeepKeys::YES) {
+                left_carrier = left_chunk.table_view();
+            } else {
+                std::vector<cudf::size_type> to_keep;
+                std::ranges::copy_if(
+                    std::ranges::iota_view(0, left_chunk.table_view().num_columns()),
+                    std::back_inserter(to_keep),
+                    [&](auto i) { return std::ranges::find(left_on, i) == left_on.end(); }
+                );
+                left_carrier = left_chunk.table_view().select(to_keep);
+            }
+            co_await ch_out->send(semi_join_chunk(
+                ctx,
+                left_chunk,
+                std::move(right_chunk),
+                left_carrier,
+                left_on,
+                right_on,
+                left_msg.sequence_number(),
+                &left_event
+            ));
+        }
+    }
     streaming::Node shuffle(
         std::shared_ptr<streaming::Context> ctx,
         std::shared_ptr<streaming::Channel> ch_in,
@@ Expand All / @@ -360,6 +555,7 @@ streaming::Node shuffle( @@
         while (true) {
             auto msg = co_await ch_in->receive();
             if (msg.empty()) {
+                ctx->comm()->logger().debug("Shuffle: no more input");
                 break;
             }
             auto chunk = co_await msg.release<streaming::TableChunk>().make_available(ctx);
@@ Expand Down @@

cpp/benchmarks/streaming/ndsh/join.hpp

-Original file line number
+Diff line change
@@ -1,5 +1,5 @@
     /**
-     * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+     * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
      * SPDX-License-Identifier: Apache-2.0
      */
@@ Expand Down Expand Up / @@ -120,6 +120,67 @@ enum class KeepKeys : bool { @@
         KeepKeys keep_keys = KeepKeys::YES
     );
+    /**
+     * @brief Perform a streaming left semi join between two tables.
+     *
+     * @note This performs a broadcast join, broadcasting the table represented by the `left`
+     * channel to all ranks, and then streaming through the chunks of the `right` channel.
+     * The `right` channel is required to provide hash-partitioned data in-order.
+     * All of the chunks from the `left` channel must fit in memory at once.
+     *
+     * @param ctx Streaming context.
+     * @param left Channel of `TableChunk`s.
+     * @param right Channel of `TableChunk`s in hash-partitioned order (shuffled).
+     * @param ch_out Output channel of `TableChunk`s.
+     * @param left_on Column indices of the keys in the left table.
+     * @param right_on Column indices of the keys in the right table.
+     * @param tag Disambiguating tag for the broadcast of the left table.
+     * @param keep_keys Does the result contain the key columns, or only "carrier" value
+     * columns
+     *
+     * @return Coroutine representing the completion of the join.
+     */
+    streaming::Node left_semi_join_broadcast_left(
+        std::shared_ptr<streaming::Context> ctx,
+        // We will always choose left as build table and do "broadcast" joins
+        std::shared_ptr<streaming::Channel> left,
+        std::shared_ptr<streaming::Channel> right,
+        std::shared_ptr<streaming::Channel> ch_out,
+        std::vector<cudf::size_type> left_on,
+        std::vector<cudf::size_type> right_on,
+        OpID tag,
+        KeepKeys keep_keys
+    );
+    /**
+     * @brief Perform a streaming left semi join between two tables.
+     *
+     * @note This performs a shuffle join, the left and right channels are required to provide
+     * hash-partitioned data in-order.
+     *
+     * @param ctx Streaming context.
+     * @param left Channel of `TableChunk`s in hash-partitioned order.
+     * @param right Channel of `TableChunk`s in matching hash-partitioned order.
+     * @param ch_out Output channel of `TableChunk`s.
+     * @param left_on Column indices of the keys in the left table.
+     * @param right_on Column indices of the keys in the right table.
+     * @param tag Disambiguating tag for the broadcast of the left table.
+     * @param keep_keys Does the result contain the key columns, or only "carrier" value
+     * columns
+     *
+     * @return Coroutine representing the completion of the join.
+     */
+    streaming::Node left_semi_join_shuffle(
+        std::shared_ptr<streaming::Context> ctx,
+        std::shared_ptr<streaming::Channel> left,
+        std::shared_ptr<streaming::Channel> right,
+        std::shared_ptr<streaming::Channel> ch_out,
+        std::vector<cudf::size_type> left_on,
+        std::vector<cudf::size_type> right_on,
+        KeepKeys keep_keys = KeepKeys::YES
+    );
     /**
      * @brief Shuffle the input channel by hash-partitioning on given key columns.
      *
@@ Expand Down @@

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Streaming Q4 implementation #710

Uh oh!

Diff view

Diff view

There are no files selected for viewing

wence- Jan 15, 2026

Uh oh!

TomAugspurger Jan 29, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Streaming Q4 implementation #710

Are you sure you want to change the base?

Uh oh!

Streaming Q4 implementation #710

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

wence- Jan 15, 2026

Choose a reason for hiding this comment

Uh oh!

TomAugspurger Jan 29, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!