Update cpp/src/arrow/tensor_benchmark.cc

Co-authored-by: Joris Van den Bossche <[email protected]>
apache · Mar 22, 2024 · 8f7fc4f · 8f7fc4f
1 parent 2e1e958
commit 8f7fc4f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/cpp/src/arrow/tensor_benchmark.cc b/cpp/src/arrow/tensor_benchmark.cc
@@ -30,7 +30,8 @@ static void BatchToTensorSimple(benchmark::State& state) {
   using CType = typename ValueType::c_type;
   std::shared_ptr<DataType> ty = TypeTraits<ValueType>::type_singleton();
 
-  const int64_t kNumRows = state.range(0) / sizeof(CType);
+  const int64_t num_cols = state.range(1);
+  const int64_t num_rows = state.range(0) / num_cols / sizeof(CType);
   arrow::random::RandomArrayGenerator gen_{42};
 
   std::vector<std::shared_ptr<Field>> fields = {};