|
1 | 1 | /* |
2 | | - * Copyright (c) 2025, NVIDIA CORPORATION. |
| 2 | + * Copyright (c) 2025-2026, NVIDIA CORPORATION. |
3 | 3 | * |
4 | 4 | * Licensed under the Apache License, Version 2.0 (the "License"); |
5 | 5 | * you may not use this file except in compliance with the License. |
@@ -867,26 +867,24 @@ TEST_F(ShuffleSplitTests, MixedValidity) |
867 | 867 | cudf::table_view expected_t{{static_cast<cudf::column_view>(*expected)}}; |
868 | 868 |
|
869 | 869 | // make the concatenated shuffle_split partitions |
870 | | - std::vector< |
871 | | - std::pair<spark_rapids_jni::shuffle_split_result, spark_rapids_jni::shuffle_split_metadata>> |
872 | | - shuf; |
| 870 | + std::vector<spark_rapids_jni::shuffle_split_output> shuf; |
873 | 871 | size_t total_size = 0; |
874 | 872 | for (size_t idx = 0; idx < partition_views.size(); idx++) { |
875 | 873 | shuf.push_back(spark_rapids_jni::shuffle_split( |
876 | 874 | cudf::table_view{{partition_views[idx]}}, {}, stream, mr)); |
877 | | - total_size += shuf.back().first.partitions->size(); |
| 875 | + total_size += shuf.back().result.partitions->size(); |
878 | 876 | } |
879 | 877 | rmm::device_uvector<uint8_t> full{total_size, stream, mr}; |
880 | 878 | rmm::device_uvector<size_t> full_offsets{partition_views.size() + 1, stream, mr}; |
881 | 879 | std::vector<size_t> h_full_offsets(partition_views.size() + 1); |
882 | 880 | size_t pos = 0; |
883 | 881 | for (size_t idx = 0; idx < partition_views.size(); idx++) { |
884 | 882 | cudaMemcpy(static_cast<uint8_t*>(full.data()) + pos, |
885 | | - shuf[idx].first.partitions->data(), |
886 | | - shuf[idx].first.partitions->size(), |
| 883 | + shuf[idx].result.partitions->data(), |
| 884 | + shuf[idx].result.partitions->size(), |
887 | 885 | cudaMemcpyDeviceToDevice); |
888 | 886 | h_full_offsets[idx] = pos; |
889 | | - pos += shuf[idx].first.partitions->size(); |
| 887 | + pos += shuf[idx].result.partitions->size(); |
890 | 888 | } |
891 | 889 | h_full_offsets[partition_views.size()] = pos; |
892 | 890 | cudaMemcpy(full_offsets.data(), |
|
0 commit comments