You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by GitBox <gi...@apache.org> on 2022/03/02 00:00:52 UTC
[GitHub] [arrow] save-buffer opened a new pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
save-buffer opened a new pull request #12537:
URL: https://github.com/apache/arrow/pull/12537
This PR contains an implementation of a multithreaded TPC-H dbgen, as well as an implementation of Q1 as a google benchmark. The advantage of this dbgen approach is that it is a scan node: it generates data on the fly and streams it over. As a result, I was for instance able to run scale factor 1000 on Q1 on my desktop with only 32 GB of RAM.
I did verify results of Q1. They don't exactly match the reference results, but they are quite close and well within what I'd expect the variance to be between random number generators.
```
-------------------------------------------------------------
Benchmark Time CPU Iterations
-------------------------------------------------------------
BM_Tpch_Q1/SF:1 186609936 ns 268825 ns 100
BM_Tpch_Q1/SF:10 1858114140 ns 276741 ns 10
BM_Tpch_Q1/SF:100 18561088470 ns 273067 ns 1
BM_Tpch_Q1/SF:1000 186103719755 ns 289445 ns 1
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1057478174
From the R side, yeah. The hard part is that we want to take each of these Record Batch Readers and write them to parquet files, but we don't have an object in the R bindings that is designed to hold multiple record batch readers and kick off the writing for both.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1057480303
Though we _should_ be able to construct a plan that includes write nodes for both of those recordbatchreaders in C++ and then let C++ handle the rest... (no need to try and shoe horn that multithreading into R which will be reluctant to make that easy!)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r818829466
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3758 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ random::pcg32_fast rng_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ int scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_generated_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
Review comment:
If you really care about speed here (why do you?), then perhaps you can instead use bisection (which can also be hardcoded as a series of if statements).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r824986313
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string> &columns,
+ const std::vector<std::shared_ptr<DataType>> &types,
+ const std::unordered_map<std::string, int> &name_map,
+ std::vector<int> &gen_list)
+ {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if(columns.empty())
+ {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for(auto pair : name_map)
+ {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ }
+ else
+ {
+ for(const std::string &col : columns)
+ {
+ auto entry = name_map.find(col);
+ if(entry == name_map.end())
+ return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+ }
+
+ static TpchText g_text;
+
+ Status TpchText::InitIfNeeded(random::pcg32_fast &rng)
+ {
+ if(done_.load())
+ return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(!text_)
+ {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char *out = reinterpret_cast<char *>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+ while(done_.load() == false)
+ {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while(GenerateSentence(try_offset, rng, temp_buff))
+ known_valid_offset = try_offset;
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(done_.load())
+ return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ int64_t memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ std::memcpy(out + generated_offset_, temp_buff, memcpy_size);
+ generated_offset_ += memcpy_size;
+ if(generated_offset_ == kTextBytes)
+ done_.store(true);
+ }
Review comment:
```suggestion
bool finishes = false;
{
std::lock_guard<std::mutex> lock(text_guard_);
if(done_.load())
return Status::OK();
int64_t bytes_remaining = kTextBytes - generated_offset_;
int64_t memcpy_size = std::min(known_valid_offset, bytes_remaining);
int64_t start = out + generated_offset_;
generated_offset_ += memcpy_size;
if(generated_offset_ == kTextBytes)
finishes = true;
}
std::memcpy(start, temp_buff, memcpy_size);
if (finishes) {
done_.store(true);
}
```
I doubt this would have a significant improvement to performance but it might be slightly faster.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1065137496
I just rebased to avoid conflicts in the R C++, so please pull that @save-buffer (sorry if you've already started looking at these and that's a pain! I can deal with conflicts in r/ if you need me to)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r825009867
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string> &columns,
+ const std::vector<std::shared_ptr<DataType>> &types,
+ const std::unordered_map<std::string, int> &name_map,
+ std::vector<int> &gen_list)
+ {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if(columns.empty())
+ {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for(auto pair : name_map)
+ {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ }
+ else
+ {
+ for(const std::string &col : columns)
+ {
+ auto entry = name_map.find(col);
+ if(entry == name_map.end())
+ return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+ }
+
+ static TpchText g_text;
+
+ Status TpchText::InitIfNeeded(random::pcg32_fast &rng)
+ {
+ if(done_.load())
+ return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(!text_)
+ {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char *out = reinterpret_cast<char *>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+ while(done_.load() == false)
+ {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while(GenerateSentence(try_offset, rng, temp_buff))
+ known_valid_offset = try_offset;
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(done_.load())
+ return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ int64_t memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ std::memcpy(out + generated_offset_, temp_buff, memcpy_size);
+ generated_offset_ += memcpy_size;
+ if(generated_offset_ == kTextBytes)
+ done_.store(true);
+ }
+ }
+ return Status::OK();
+ }
+
+ Result<Datum> TpchText::GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng)
+ {
+ RETURN_NOT_OK(InitIfNeeded(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer, AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for(size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + length_dist(rng);
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer, AllocateBuffer(offsets[num_comments]));
+ char *comments = reinterpret_cast<char *>(comment_buffer->mutable_data());
+ for(size_t i = 0; i < num_comments; i++)
+ {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments, { nullptr, std::move(offset_buffer), std::move(comment_buffer) });
+ return std::move(ad);
+ }
+
+ Result<Datum> RandomVString(
+ random::pcg32_fast &rng,
+ int64_t num_rows,
+ int32_t min_length,
+ int32_t max_length)
+ {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff, AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for(int64_t i = 1; i <= num_rows; i++)
+ offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff, AllocateBuffer(offsets[num_rows]));
+ char *str = reinterpret_cast<char *>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for(int32_t i = 0; i < offsets[num_rows]; i++)
+ str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, { nullptr, std::move(offset_buff), std::move(str_buff) });
+ return std::move(ad);
+ }
+
+ void AppendNumber(char *&out, int num_digits, int32_t x)
+ {
+ out += (num_digits - 1);
+ while(x > 0)
+ {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+ }
+
+ void GeneratePhoneNumber(
+ char *out,
+ random::pcg32_fast &rng,
+ int32_t country)
+ {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+ }
+
+ static constexpr uint32_t STARTDATE = 8035; // January 1, 1992 is 8035 days after January 1, 1970
+ static constexpr uint32_t CURRENTDATE = 9298; // June 17, 1995 is 9298 days after January 1, 1970
+ static constexpr uint32_t ENDDATE = 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+ const char *NameParts[] =
+ {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque", "black", "blanched", "blue",
+ "blush", "brown", "burlywood", "burnished", "chartreuse", "chiffon", "chocolate", "coral",
+ "cornflower", "cornsilk", "cream", "cyan", "dark", "deep", "dim", "dodger", "drab", "firebrick",
+ "floral", "forest", "frosted", "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender", "lawn", "lemon", "light", "lime", "linen",
+ "magenta", "maroon", "medium", "metallic", "midnight", "mint", "misty", "moccasin", "navajo",
+ "navy", "olive", "orange", "orchid", "pale", "papaya", "peach", "peru", "pink", "plum", "powder",
+ "puff", "purple", "red", "rose", "rosy", "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel", "tan", "thistle", "tomato", "turquoise", "violet",
+ "wheat", "white", "yellow",
+ };
+ static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+ const char *Types_1[] =
+ {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+ };
+ static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+ const char *Types_2[] =
+ {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+ };
+ static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+ const char *Types_3[] =
+ {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+ };
+ static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+ const char *Containers_1[] =
+ {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+ };
+ static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+ const char *Containers_2[] =
+ {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+ };
+ static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+ const char *Segments[] =
+ {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+ };
+ static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+ const char *Priorities[] =
+ {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+ };
+ static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+ const char *Instructions[] =
+ {
+ "DELIVER IN PERSON", "COLLECT COD", "NONE", "TAKE BACK RETURN",
+ };
+ static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+ const char *Modes[] =
+ {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+ };
+ static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+ const char *Nouns[] =
+ {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ", "dependencies ", "excuses ",
+ "platelets ", "asymptotes ", "courts ", "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ", "forges ", "braids ", "hockey players ", "frays ",
+ "warhorses ", "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ", "waters ", "orbits ", "gifts ", "sheaves ",
+ "depths ", "sentiments ", "decoys ", "realms ", "pains ", "grouches ", "escapades ",
+ };
+ static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+ const char *Verbs[] =
+ {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ", "boost ", "affix ", "detect ", "integrate ",
+ "maintain ", "nod ", "was ", "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ", "poach ", "serve ", "run ", "dazzle ",
+ "snooze ", "doze ", "unwind ", "kindle ", "play ", "hang ", "believe ", "doubt ",
+ };
+ static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+ const char *Adjectives[] =
+ {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ", "quiet ", "ruthless ", "thin ",
+ "close ", "dogged ", "daring ", "brave ", "stealthy ", "permanent ", "enticing ", "idle ", "busy ",
+ "regular ", "final ", "ironic ", "even ", "bold ", "silent ",
+ };
+ static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+ const char *Adverbs[] =
+ {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ", "blithely ", "quickly ", "fluffily ",
+ "slowly ", "quietly ", "ruthlessly ", "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ", "ironically ", "evenly ", "boldly ",
+ "silently ",
+ };
+ static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+ const char *Prepositions[] =
+ {
+ "about ", "above ", "according to ", "across ", "after ", "against ", "along ", "alongside of ", "among ",
+ "around ", "at ", "atop ", "before ", "behind ", "beneath ", "beside ", "besides ", "between ", "beyond ",
+ "beyond ", "by ", "despite ", "during ", "except ", "for ", "from ", "in place of ", "inside ", "instead of ",
+ "into ", "near ", "of ", "on ", "outside ", "over ", "past ", "since ", "through ", "throughout ", "to ",
+ "toward ", "under ", "until ", "up ", "upon ", "without ", "with ", "within ",
+ };
+ static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+ const char *Auxiliaries[] =
+ {
+ "do ", "may ", "might ", "shall ", "will ", "would ", "can ", "could ", "should ", "ought to ", "must ",
+ "will have to ", "shall have to ", "could have to ", "should have to ", "must have to ", "need to ", "try to ",
+ };
+ static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+ const char *Terminators[] =
+ {
+ ".", ";", ":", "?", "!", "--",
+ };
+ static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+ bool TpchText::GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices)
+ {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char *word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if(offset + length > kChunkSize)
+ return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+ }
+
+ bool TpchText::GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+ }
+
+ bool TpchText::GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+ }
+
+ bool TpchText::GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+ }
+
+ bool TpchText::GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+ }
+
+ bool TpchText::GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+ }
+
+ bool TpchText::GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+ }
+
+ bool TpchText::GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if(result)
+ std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+ }
+
+ bool TpchText::GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char *comma_space = ", ";
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ GenerateAdverb(offset, rng, arr);
+ GenerateAdjective(offset, rng, arr);
+ GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ bool TpchText::GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ bool TpchText::GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ const char *the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+ }
+
+ bool TpchText::GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 5:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ using GenerateColumnFn = std::function<Status(size_t)>;
+ class PartAndPartSupplierGenerator
+ {
+ public:
+ Status Init(
+ size_t num_threads,
+ int64_t batch_size,
+ float scale_factor)
+ {
+ if(!inited_)
+ {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ thread_local_data_.resize(num_threads);
+ for(ThreadLocalData &tld : thread_local_data_)
+ {
+ // 5 is the maximum number of different strings we need to concatenate
+ tld.string_indices.resize(5 * batch_size_);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const
+ {
+ return part_batches_generated_.load();
+ }
+
+ int64_t partsupp_batches_generated() const
+ {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, part_types_, part_name_map_, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, partsupp_types_, partsupp_name_map_, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch()
+ {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ bool all_generated = part_rows_generated_ == part_rows_to_generate_;
+ if(!part_output_queue_.empty())
+ {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ }
+ else if(all_generated)
+ {
+ return util::nullopt;
+ }
+ else
+ {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate = std::min(
+ batch_size_,
+ part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for(int col : part_cols_)
+ {
+ RETURN_NOT_OK(part_generators_[col](thread_index));
+ }
+ for(int col : partsupp_cols_)
+ RETURN_NOT_OK(partsupp_generators_[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for(size_t i = 0; i < part_cols_.size(); i++)
+ {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if(!partsupp_cols_.empty())
+ {
+ std::vector<ExecBatch> partsupp_results;
+ for(size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++)
+ {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for(size_t icol = 0; icol < partsupp_cols_.size(); icol++)
+ {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for(ExecBatch &eb : partsupp_results)
+ {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch()
+ {
+ size_t thread_index = thread_indexer_();
Review comment:
Same as above
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r825010298
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string> &columns,
+ const std::vector<std::shared_ptr<DataType>> &types,
+ const std::unordered_map<std::string, int> &name_map,
+ std::vector<int> &gen_list)
+ {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if(columns.empty())
+ {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for(auto pair : name_map)
+ {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ }
+ else
+ {
+ for(const std::string &col : columns)
+ {
+ auto entry = name_map.find(col);
+ if(entry == name_map.end())
+ return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+ }
+
+ static TpchText g_text;
+
+ Status TpchText::InitIfNeeded(random::pcg32_fast &rng)
+ {
+ if(done_.load())
+ return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(!text_)
+ {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char *out = reinterpret_cast<char *>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+ while(done_.load() == false)
+ {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while(GenerateSentence(try_offset, rng, temp_buff))
+ known_valid_offset = try_offset;
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(done_.load())
+ return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ int64_t memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ std::memcpy(out + generated_offset_, temp_buff, memcpy_size);
+ generated_offset_ += memcpy_size;
+ if(generated_offset_ == kTextBytes)
+ done_.store(true);
+ }
+ }
+ return Status::OK();
+ }
+
+ Result<Datum> TpchText::GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng)
+ {
+ RETURN_NOT_OK(InitIfNeeded(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer, AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for(size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + length_dist(rng);
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer, AllocateBuffer(offsets[num_comments]));
+ char *comments = reinterpret_cast<char *>(comment_buffer->mutable_data());
+ for(size_t i = 0; i < num_comments; i++)
+ {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments, { nullptr, std::move(offset_buffer), std::move(comment_buffer) });
+ return std::move(ad);
+ }
+
+ Result<Datum> RandomVString(
+ random::pcg32_fast &rng,
+ int64_t num_rows,
+ int32_t min_length,
+ int32_t max_length)
+ {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff, AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for(int64_t i = 1; i <= num_rows; i++)
+ offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff, AllocateBuffer(offsets[num_rows]));
+ char *str = reinterpret_cast<char *>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for(int32_t i = 0; i < offsets[num_rows]; i++)
+ str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, { nullptr, std::move(offset_buff), std::move(str_buff) });
+ return std::move(ad);
+ }
+
+ void AppendNumber(char *&out, int num_digits, int32_t x)
+ {
+ out += (num_digits - 1);
+ while(x > 0)
+ {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+ }
+
+ void GeneratePhoneNumber(
+ char *out,
+ random::pcg32_fast &rng,
+ int32_t country)
+ {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+ }
+
+ static constexpr uint32_t STARTDATE = 8035; // January 1, 1992 is 8035 days after January 1, 1970
+ static constexpr uint32_t CURRENTDATE = 9298; // June 17, 1995 is 9298 days after January 1, 1970
+ static constexpr uint32_t ENDDATE = 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+ const char *NameParts[] =
+ {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque", "black", "blanched", "blue",
+ "blush", "brown", "burlywood", "burnished", "chartreuse", "chiffon", "chocolate", "coral",
+ "cornflower", "cornsilk", "cream", "cyan", "dark", "deep", "dim", "dodger", "drab", "firebrick",
+ "floral", "forest", "frosted", "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender", "lawn", "lemon", "light", "lime", "linen",
+ "magenta", "maroon", "medium", "metallic", "midnight", "mint", "misty", "moccasin", "navajo",
+ "navy", "olive", "orange", "orchid", "pale", "papaya", "peach", "peru", "pink", "plum", "powder",
+ "puff", "purple", "red", "rose", "rosy", "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel", "tan", "thistle", "tomato", "turquoise", "violet",
+ "wheat", "white", "yellow",
+ };
+ static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+ const char *Types_1[] =
+ {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+ };
+ static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+ const char *Types_2[] =
+ {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+ };
+ static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+ const char *Types_3[] =
+ {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+ };
+ static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+ const char *Containers_1[] =
+ {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+ };
+ static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+ const char *Containers_2[] =
+ {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+ };
+ static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+ const char *Segments[] =
+ {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+ };
+ static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+ const char *Priorities[] =
+ {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+ };
+ static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+ const char *Instructions[] =
+ {
+ "DELIVER IN PERSON", "COLLECT COD", "NONE", "TAKE BACK RETURN",
+ };
+ static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+ const char *Modes[] =
+ {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+ };
+ static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+ const char *Nouns[] =
+ {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ", "dependencies ", "excuses ",
+ "platelets ", "asymptotes ", "courts ", "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ", "forges ", "braids ", "hockey players ", "frays ",
+ "warhorses ", "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ", "waters ", "orbits ", "gifts ", "sheaves ",
+ "depths ", "sentiments ", "decoys ", "realms ", "pains ", "grouches ", "escapades ",
+ };
+ static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+ const char *Verbs[] =
+ {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ", "boost ", "affix ", "detect ", "integrate ",
+ "maintain ", "nod ", "was ", "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ", "poach ", "serve ", "run ", "dazzle ",
+ "snooze ", "doze ", "unwind ", "kindle ", "play ", "hang ", "believe ", "doubt ",
+ };
+ static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+ const char *Adjectives[] =
+ {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ", "quiet ", "ruthless ", "thin ",
+ "close ", "dogged ", "daring ", "brave ", "stealthy ", "permanent ", "enticing ", "idle ", "busy ",
+ "regular ", "final ", "ironic ", "even ", "bold ", "silent ",
+ };
+ static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+ const char *Adverbs[] =
+ {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ", "blithely ", "quickly ", "fluffily ",
+ "slowly ", "quietly ", "ruthlessly ", "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ", "ironically ", "evenly ", "boldly ",
+ "silently ",
+ };
+ static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+ const char *Prepositions[] =
+ {
+ "about ", "above ", "according to ", "across ", "after ", "against ", "along ", "alongside of ", "among ",
+ "around ", "at ", "atop ", "before ", "behind ", "beneath ", "beside ", "besides ", "between ", "beyond ",
+ "beyond ", "by ", "despite ", "during ", "except ", "for ", "from ", "in place of ", "inside ", "instead of ",
+ "into ", "near ", "of ", "on ", "outside ", "over ", "past ", "since ", "through ", "throughout ", "to ",
+ "toward ", "under ", "until ", "up ", "upon ", "without ", "with ", "within ",
+ };
+ static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+ const char *Auxiliaries[] =
+ {
+ "do ", "may ", "might ", "shall ", "will ", "would ", "can ", "could ", "should ", "ought to ", "must ",
+ "will have to ", "shall have to ", "could have to ", "should have to ", "must have to ", "need to ", "try to ",
+ };
+ static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+ const char *Terminators[] =
+ {
+ ".", ";", ":", "?", "!", "--",
+ };
+ static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+ bool TpchText::GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices)
+ {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char *word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if(offset + length > kChunkSize)
+ return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+ }
+
+ bool TpchText::GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+ }
+
+ bool TpchText::GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+ }
+
+ bool TpchText::GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+ }
+
+ bool TpchText::GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+ }
+
+ bool TpchText::GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+ }
+
+ bool TpchText::GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+ }
+
+ bool TpchText::GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if(result)
+ std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+ }
+
+ bool TpchText::GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char *comma_space = ", ";
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ GenerateAdverb(offset, rng, arr);
+ GenerateAdjective(offset, rng, arr);
+ GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ bool TpchText::GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ bool TpchText::GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ const char *the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+ }
+
+ bool TpchText::GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 5:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ using GenerateColumnFn = std::function<Status(size_t)>;
+ class PartAndPartSupplierGenerator
+ {
+ public:
+ Status Init(
+ size_t num_threads,
+ int64_t batch_size,
+ float scale_factor)
+ {
+ if(!inited_)
+ {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ thread_local_data_.resize(num_threads);
+ for(ThreadLocalData &tld : thread_local_data_)
+ {
+ // 5 is the maximum number of different strings we need to concatenate
+ tld.string_indices.resize(5 * batch_size_);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const
+ {
+ return part_batches_generated_.load();
+ }
+
+ int64_t partsupp_batches_generated() const
+ {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, part_types_, part_name_map_, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, partsupp_types_, partsupp_name_map_, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch()
+ {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ bool all_generated = part_rows_generated_ == part_rows_to_generate_;
+ if(!part_output_queue_.empty())
+ {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ }
+ else if(all_generated)
+ {
+ return util::nullopt;
+ }
+ else
+ {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate = std::min(
+ batch_size_,
+ part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for(int col : part_cols_)
+ {
+ RETURN_NOT_OK(part_generators_[col](thread_index));
+ }
+ for(int col : partsupp_cols_)
+ RETURN_NOT_OK(partsupp_generators_[col](thread_index));
Review comment:
I guess I'll have to clang-format anyway before pushing :P
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] ursabot edited a comment on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
ursabot edited a comment on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082832095
Benchmark runs are scheduled for baseline = 0deefd1838bf4706232b90b72da010698010f7b6 and contender = 50fab73da58c9f76fbe675d661e56087e2f5fb52. 50fab73da58c9f76fbe675d661e56087e2f5fb52 is a master commit associated with this PR. Results will be available as each benchmark for each run completes.
Conbench compare runs links:
[Finished :arrow_down:0.0% :arrow_up:0.0%] [ec2-t3-xlarge-us-east-2](https://conbench.ursa.dev/compare/runs/3f8ac59fddef426cac31b6b2826ce5a6...e5b78c70b6fe4a699c4cb8b967fe5736/)
[Scheduled] [test-mac-arm](https://conbench.ursa.dev/compare/runs/02b74ab900d440f8b5361300f15ad15c...b870e78b617e41e98c2ade7241e0aa06/)
[Scheduled] [ursa-i9-9960x](https://conbench.ursa.dev/compare/runs/1f7740546747452aa61fece1295aa0e3...f66d8edd43724d9988628cbe4a4e4d50/)
[Scheduled] [ursa-thinkcentre-m75q](https://conbench.ursa.dev/compare/runs/c3c55d6999fe4d7eabec73cbdfa1fad6...4b9f45f6442447d3a9d758598d0e14d3/)
Supported benchmarks:
ec2-t3-xlarge-us-east-2: Supported benchmark langs: Python, R. Runs only benchmarks with cloud = True
test-mac-arm: Supported benchmark langs: C++, Python, R
ursa-i9-9960x: Supported benchmark langs: Python, R, JavaScript
ursa-thinkcentre-m75q: Supported benchmark langs: C++, Java
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829433851
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
Review comment:
actually nevermind, i think users should be able to use this class for their exec plans. this can be a provided data generator.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1070989318
> Thanks @save-buffer I agree, that will be simpler. I'm happy to take a look at tpch_dbgen_write too.
Also agree we should move this forward if the R stuff is holding it back. And would love if you could take a look at that Weston!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r824954602
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
Review comment:
Again, I need this same one to be visible across all threads.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r825178760
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string> &columns,
+ const std::vector<std::shared_ptr<DataType>> &types,
+ const std::unordered_map<std::string, int> &name_map,
+ std::vector<int> &gen_list)
+ {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if(columns.empty())
+ {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for(auto pair : name_map)
+ {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ }
+ else
+ {
+ for(const std::string &col : columns)
+ {
+ auto entry = name_map.find(col);
+ if(entry == name_map.end())
+ return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+ }
+
+ static TpchText g_text;
+
+ Status TpchText::InitIfNeeded(random::pcg32_fast &rng)
+ {
+ if(done_.load())
+ return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(!text_)
+ {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char *out = reinterpret_cast<char *>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+ while(done_.load() == false)
+ {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while(GenerateSentence(try_offset, rng, temp_buff))
+ known_valid_offset = try_offset;
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(done_.load())
+ return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ int64_t memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ std::memcpy(out + generated_offset_, temp_buff, memcpy_size);
+ generated_offset_ += memcpy_size;
+ if(generated_offset_ == kTextBytes)
+ done_.store(true);
+ }
Review comment:
Nevermind, it's the same speed between lock-free and locks (as far as I could tell, averages were the same but the lock free version had some more variance), so I just switched back to this version because it seems easier to understand.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r818046023
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,175 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/testing/future_util.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/compute/cast.h"
+
+namespace arrow
+{
+namespace compute
+{
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>> &sink_gen, int scale_factor)
Review comment:
Could we make scale_factor a float? DBgen itself actually can generate scale_factors of 0.1, 0.01 (and we use those for testing _a lot_)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829615634
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
Review comment:
Agreed. We can take a seed in the constructor. If not specified I think defaulting to non-deterministic is ok.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829317347
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ if (!partsupp_output_queue_.empty()) {
+ ExecBatch result = std::move(partsupp_output_queue_.front());
+ partsupp_output_queue_.pop();
+ return std::move(result);
+ }
+ }
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+ if (!part_cols_.empty()) {
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch part_batch,
+ ExecBatch::Make(std::move(part_result)));
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ part_output_queue_.emplace(std::move(part_batch));
+ }
+ }
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ for (size_t i = 1; i < partsupp_results.size(); i++)
+ partsupp_output_queue_.emplace(std::move(partsupp_results[i]));
+ }
+ return std::move(partsupp_results[0]);
+ }
+
+ private:
+#define FOR_EACH_PART_COLUMN(F) \
+ F(P_PARTKEY) \
+ F(P_NAME) \
+ F(P_MFGR) \
+ F(P_BRAND) \
+ F(P_TYPE) \
+ F(P_SIZE) \
+ F(P_CONTAINER) \
+ F(P_RETAILPRICE) \
+ F(P_COMMENT)
+
+#define FOR_EACH_PARTSUPP_COLUMN(F) \
+ F(PS_PARTKEY) \
+ F(PS_SUPPKEY) \
+ F(PS_AVAILQTY) \
+ F(PS_SUPPLYCOST) \
+ F(PS_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct PART {
+ enum {
+ FOR_EACH_PART_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct PARTSUPP {
+ enum {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, PART::col},
+ const std::unordered_map<std::string, int> kPartNameMap = {
+ FOR_EACH_PART_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, PARTSUPP::col},
+ const std::unordered_map<std::string, int> kPartsuppNameMap = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kPartGenerators = {FOR_EACH_PART_COLUMN(MAKE_FN_ARRAY)};
+ std::vector<GenerateColumnFn> kPartsuppGenerators = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kPartTypes = {
+ int32(), utf8(), fixed_size_binary(25), fixed_size_binary(10),
+ utf8(), int32(), fixed_size_binary(10), decimal(12, 2),
+ utf8(),
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kPartsuppTypes = {
+ int32(), int32(), int32(), decimal(12, 2), utf8(),
+ };
+
+ Status AllocatePartBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.part[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.part_to_generate * byte_width));
+ ArrayData ad(kPartTypes[column], tld.part_to_generate, {nullptr, std::move(buff)});
+ tld.part[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status P_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_PARTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_PARTKEY));
+ int32_t* p_partkey = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) {
+ p_partkey[i] = static_cast<int32_t>(tld.partkey_start + i + 1);
+ ARROW_DCHECK(1 <= p_partkey[i] && p_partkey[i] <= part_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_NAME].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(
+ 0, static_cast<uint8_t>(kNumNameParts - 1));
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dist(tld.rng));
+ tld.string_indices[irow * 5 + ipart] = name_part_index;
+ string_length += std::strlen(NameParts[name_part_index]);
+ }
+ // Add 4 because there is a space between each word (i.e. four spaces)
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length + 4);
+ }
+ // Add an extra byte for the space after in the very last string.
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate] + 1));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 5 + ipart];
+ const char* part = NameParts[name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ *row++ = ' ';
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_NAME], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_NAME] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_MFGR(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_MFGR].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* manufacturer = "Manufacturer#";
+ const size_t manufacturer_length = std::strlen(manufacturer);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_MFGR));
+ char* p_mfgr = reinterpret_cast<char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ std::strncpy(p_mfgr + byte_width * irow, manufacturer, byte_width);
+ char mfgr_id = '0' + dist(tld.rng);
+ *(p_mfgr + byte_width * irow + manufacturer_length) = mfgr_id;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_BRAND(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_BRAND].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_MFGR(thread_index));
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* brand = "Brand#";
+ const size_t brand_length = std::strlen(brand);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_BRAND));
+ const char* p_mfgr = reinterpret_cast<const char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->data());
+ char* p_brand = reinterpret_cast<char*>(
+ tld.part[PART::P_BRAND].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_BRAND]);
+ int32_t mfgr_byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ const size_t mfgr_id_offset = std::strlen("Manufacturer#");
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = p_brand + byte_width * irow;
+ char mfgr_id = *(p_mfgr + irow * mfgr_byte_width + mfgr_id_offset);
+ char brand_id = '0' + dist(tld.rng);
+ std::strncpy(row, brand, byte_width);
+ *(row + brand_length) = mfgr_id;
+ *(row + brand_length + 1) = brand_id;
+ irow += 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_TYPE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_TYPE].kind() == Datum::NONE) {
+ using D = std::uniform_int_distribution<uint32_t>;
+ D dists[] = {
+ D{0, static_cast<uint8_t>(kNumTypes_1 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_2 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_3 - 1)},
+ };
+
+ const char** types[] = {Types_1, Types_2, Types_3};
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dists[ipart](tld.rng));
+ tld.string_indices[irow * 3 + ipart] = name_part_index;
+ string_length += std::strlen(types[ipart][name_part_index]);
+ }
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate]));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 3 + ipart];
+ const char* part = types[ipart][name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_TYPE], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_TYPE] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_SIZE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_SIZE].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(1, 50);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_SIZE));
+ int32_t* p_size = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_SIZE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) p_size[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status P_CONTAINER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_CONTAINER].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist1(
+ 0, static_cast<uint8_t>(kNumContainers_1 - 1));
+ std::uniform_int_distribution<int> dist2(
+ 0, static_cast<uint8_t>(kNumContainers_2 - 1));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_CONTAINER));
+ char* p_container = reinterpret_cast<char*>(
+ tld.part[PART::P_CONTAINER].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_CONTAINER]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int container1_idx = dist1(tld.rng);
+ int container2_idx = dist2(tld.rng);
+ const char* container1 = Containers_1[container1_idx];
+ const char* container2 = Containers_2[container2_idx];
+ size_t container1_length = std::strlen(container1);
+ size_t container2_length = std::strlen(container2);
+
+ char* row = p_container + byte_width * irow;
+ std::strncpy(row, container1, byte_width);
+ std::memcpy(row + container1_length, container2, container2_length);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_RETAILPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_RETAILPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_RETAILPRICE));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+ Decimal128* p_retailprice = reinterpret_cast<Decimal128*>(
+ tld.part[PART::P_RETAILPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int32_t partkey = p_partkey[irow];
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ p_retailprice[irow] = {retail_price};
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.part[PART::P_COMMENT],
+ g_text.GenerateComments(tld.part_to_generate, 5, 22, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ int64_t PartsuppBatchesToGenerate(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t num_batches = (ps_to_generate + batch_size_ - 1) / batch_size_;
+ return num_batches;
+ }
+
+ Status InitPartsupp(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.generated_partsupp.reset();
+ int64_t num_batches = PartsuppBatchesToGenerate(thread_index);
+ tld.partsupp.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.partsupp) {
+ batch.resize(PARTSUPP::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocatePartSuppBatch(size_t thread_index, size_t ibatch, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartsuppTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kPartsuppTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.partsupp[ibatch][column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status PS_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_PARTKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_PARTKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_PARTKEY));
+ int32_t* ps_partkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run; ipartsupp++, irun++)
+ ps_partkey[batch_offset++] = p_partkey[ipart];
+
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPKEY));
+ int32_t* ps_suppkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run;
+ ipartsupp++, irun++) {
+ int32_t supplier = static_cast<int32_t>(ipartsupp);
+ int32_t partkey = p_partkey[ipart];
+ ps_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_AVAILQTY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_AVAILQTY]) {
+ tld.generated_partsupp[PARTSUPP::PS_AVAILQTY] = true;
+ std::uniform_int_distribution<int32_t> dist(1, 9999);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_AVAILQTY));
+ int32_t* ps_availqty =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for (int64_t irun = 0; irun < next_run; irun++) ps_availqty[irun] = dist(tld.rng);
+
+ tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPLYCOST(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST] = true;
+ std::uniform_int_distribution<int64_t> dist(100, 100000);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(
+ AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPLYCOST));
+ Decimal128* ps_supplycost =
+ reinterpret_cast<Decimal128*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for (int64_t irun = 0; irun < next_run; irun++)
+ ps_supplycost[irun] = {dist(tld.rng)};
+
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PARTSUPP::PS_COMMENT].kind() == Datum::NONE) {
+ int64_t irow = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ int64_t num_rows = std::min(batch_size_, ps_to_generate - irow);
+ ARROW_ASSIGN_OR_RAISE(tld.partsupp[ibatch][PARTSUPP::PS_COMMENT],
+ g_text.GenerateComments(num_rows, 49, 198, tld.rng));
+ irow += num_rows;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ std::vector<Datum> part;
+ std::vector<int8_t> string_indices;
+ int64_t part_to_generate{0};
+ int64_t partkey_start{0};
+
+ std::vector<std::vector<Datum>> partsupp;
+ std::bitset<PARTSUPP::kNumCols> generated_partsupp;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex part_output_queue_mutex_;
+ std::mutex partsupp_output_queue_mutex_;
+ std::queue<ExecBatch> part_output_queue_;
+ std::queue<ExecBatch> partsupp_output_queue_;
+ int64_t batch_size_{0};
+ float scale_factor_{0};
+ int64_t part_rows_to_generate_{0};
+ int64_t part_rows_generated_{0};
+ std::vector<int> part_cols_;
+ std::vector<int> partsupp_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<int64_t> part_batches_generated_ = {0};
+ std::atomic<int64_t> partsupp_batches_generated_ = {0};
+ static constexpr int64_t kPartSuppRowsPerPart = 4;
+};
+
+class OrdersAndLineItemGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ tld.items_per_order.resize(batch_size_);
+ tld.rng.seed(seq);
+ }
+ orders_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000 * 10);
+ }
+ return Status::OK();
+ }
+
+ int64_t orders_batches_generated() const { return orders_batches_generated_.load(); }
+
+ int64_t lineitem_batches_generated() const {
+ return lineitem_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetOrdersOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kOrdersTypes, kOrdersNameMap, orders_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetLineItemOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kLineitemTypes, kLineitemNameMap, lineitem_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextOrdersBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if (!orders_output_queue_.empty()) {
+ ExecBatch batch = std::move(orders_output_queue_.front());
+ orders_output_queue_.pop();
+ return std::move(batch);
+ } else if (orders_rows_generated_ == orders_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate =
+ std::min(batch_size_, orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.first_batch_offset = 0;
+ tld.generated_lineitem.reset();
+
+ for (int col : orders_cols_) RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for (int col : lineitem_cols_) RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for (size_t i = 0; i < orders_cols_.size(); i++) {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ if (!lineitem_cols_.empty()) {
+ std::vector<ExecBatch> lineitem_results;
+ for (size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++) {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for (size_t icol = 0; icol < lineitem_cols_.size(); icol++) {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(lineitem_output_queue_mutex_);
+ for (ExecBatch& eb : lineitem_results) {
+ lineitem_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(orders_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextLineItemBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ExecBatch queued;
+ bool from_queue = false;
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ if (!lineitem_output_queue_.empty()) {
+ queued = std::move(lineitem_output_queue_.front());
+ lineitem_output_queue_.pop();
+ from_queue = true;
+ }
+ }
+ tld.first_batch_offset = 0;
+ if (from_queue) {
+ ARROW_DCHECK(queued.length <= batch_size_);
+ tld.first_batch_offset = queued.length;
+ if (queued.length == batch_size_) return std::move(queued);
+ }
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if (orders_rows_generated_ == orders_rows_to_generate_) {
+ if (from_queue) return std::move(queued);
+ return util::nullopt;
+ }
+
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate =
+ std::min(batch_size_, orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1ll);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.generated_lineitem.reset();
+ if (from_queue) {
+ lineitem_batches_generated_.fetch_sub(1);
+ for (size_t i = 0; i < lineitem_cols_.size(); i++)
+ if (tld.lineitem[0][lineitem_cols_[i]].kind() == Datum::NONE)
+ tld.lineitem[0][lineitem_cols_[i]] = std::move(queued[i]);
+ }
+
+ for (int col : orders_cols_) RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for (int col : lineitem_cols_) RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ if (!orders_cols_.empty()) {
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for (size_t i = 0; i < orders_cols_.size(); i++) {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch orders_batch,
+ ExecBatch::Make(std::move(orders_result)));
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ orders_output_queue_.emplace(std::move(orders_batch));
+ }
+ }
+ std::vector<ExecBatch> lineitem_results;
+ for (size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++) {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for (size_t icol = 0; icol < lineitem_cols_.size(); icol++) {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ lineitem_batches_generated_.fetch_add(static_cast<int64_t>(lineitem_results.size()));
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ for (size_t i = 1; i < lineitem_results.size(); i++)
+ lineitem_output_queue_.emplace(std::move(lineitem_results[i]));
+ }
+ return std::move(lineitem_results[0]);
+ }
+
+ private:
+#define FOR_EACH_ORDERS_COLUMN(F) \
+ F(O_ORDERKEY) \
+ F(O_CUSTKEY) \
+ F(O_ORDERSTATUS) \
+ F(O_TOTALPRICE) \
+ F(O_ORDERDATE) \
+ F(O_ORDERPRIORITY) \
+ F(O_CLERK) \
+ F(O_SHIPPRIORITY) \
+ F(O_COMMENT)
+
+#define FOR_EACH_LINEITEM_COLUMN(F) \
+ F(L_ORDERKEY) \
+ F(L_PARTKEY) \
+ F(L_SUPPKEY) \
+ F(L_LINENUMBER) \
+ F(L_QUANTITY) \
+ F(L_EXTENDEDPRICE) \
+ F(L_DISCOUNT) \
+ F(L_TAX) \
+ F(L_RETURNFLAG) \
+ F(L_LINESTATUS) \
+ F(L_SHIPDATE) \
+ F(L_COMMITDATE) \
+ F(L_RECEIPTDATE) \
+ F(L_SHIPINSTRUCT) \
+ F(L_SHIPMODE) \
+ F(L_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct ORDERS {
+ enum {
+ FOR_EACH_ORDERS_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct LINEITEM {
+ enum {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, ORDERS::col},
+ const std::unordered_map<std::string, int> kOrdersNameMap = {
+ FOR_EACH_ORDERS_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, LINEITEM::col},
+ const std::unordered_map<std::string, int> kLineitemNameMap = {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ const std::vector<GenerateColumnFn> kOrdersGenerators = {
+ FOR_EACH_ORDERS_COLUMN(MAKE_FN_ARRAY)};
+ const std::vector<GenerateColumnFn> kLineitemGenerators = {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kOrdersTypes = {int32(),
+ int32(),
+ fixed_size_binary(1),
+ decimal(12, 2),
+ date32(),
+ fixed_size_binary(15),
+ fixed_size_binary(15),
+ int32(),
+ utf8()};
+
+ const std::vector<std::shared_ptr<DataType>> kLineitemTypes = {
+ int32(),
+ int32(),
+ int32(),
+ int32(),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ fixed_size_binary(1),
+ fixed_size_binary(1),
+ date32(),
+ date32(),
+ date32(),
+ fixed_size_binary(25),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status AllocateOrdersBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.orders[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.orders_to_generate * byte_width));
+ ArrayData ad(kOrdersTypes[column], tld.orders_to_generate,
+ {nullptr, std::move(buff)});
+ tld.orders[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status O_ORDERKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERKEY));
+ int32_t* o_orderkey = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ int32_t orderkey_index = static_cast<int32_t>(tld.orderkey_start + i);
+ int32_t index_of_run = orderkey_index / 8;
+ int32_t index_in_run = orderkey_index % 8;
+ o_orderkey[i] = (index_of_run * 32 + index_in_run + 1);
+ ARROW_DCHECK(1 <= o_orderkey[i] && o_orderkey[i] <= 4 * orders_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CUSTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_CUSTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CUSTKEY));
+
+ // Spec says it must be a random number between 1 and SF*150000 that is not
+ // divisible by 3. Rather than repeatedly generating numbers until we get to
+ // a non-divisible-by-3 number, we just generate a number between
+ // 0 and SF * 50000 - 1, multiply by 3, and then add either 1 or 2.
+ int32_t sf_50k = static_cast<int32_t>(scale_factor_ * 50000);
+ std::uniform_int_distribution<int32_t> base_dist(0, sf_50k - 1);
+ std::uniform_int_distribution<int32_t> offset_dist(1, 2);
+ int32_t* o_custkey = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_CUSTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++)
+ o_custkey[i] = 3 * base_dist(tld.rng) + offset_dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERSTATUS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERSTATUS].kind() == Datum::NONE) {
+ RETURN_NOT_OK(L_LINESTATUS(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERSTATUS));
+
+ char* o_orderstatus = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_ORDERSTATUS].array()->buffers[1]->mutable_data());
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ bool all_f = true;
+ bool all_o = true;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ const char* l_linestatus = reinterpret_cast<const char*>(
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->buffers[1]->data());
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++, batch_offset++) {
+ all_f &= l_linestatus[batch_offset] == 'F';
+ all_o &= l_linestatus[batch_offset] == 'O';
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ ARROW_DCHECK(!(all_f && all_o));
+ if (all_f)
+ o_orderstatus[iorder] = 'F';
+ else if (all_o)
+ o_orderstatus[iorder] = 'O';
+ else
+ o_orderstatus[iorder] = 'P';
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_TOTALPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_TOTALPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(L_EXTENDEDPRICE(thread_index));
+ RETURN_NOT_OK(L_TAX(thread_index));
+ RETURN_NOT_OK(L_DISCOUNT(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_TOTALPRICE));
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ int64_t sum = 0;
+ Decimal128* o_totalprice = reinterpret_cast<Decimal128*>(
+ tld.orders[ORDERS::O_TOTALPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ const Decimal128* l_extendedprice = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->buffers[1]->data());
+ const Decimal128* l_tax = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->data());
+ const Decimal128* l_discount = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->buffers[1]->data());
+
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++, batch_offset++) {
+ int64_t eprice = static_cast<int64_t>(l_extendedprice[batch_offset]);
+ int64_t tax = static_cast<int64_t>(l_tax[batch_offset]);
+ int64_t discount = static_cast<int64_t>(l_discount[batch_offset]);
+ sum += (eprice * (100 + tax) * (100 - discount));
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ sum /= 100 * 100;
+ o_totalprice[iorder] = {sum};
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERDATE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERDATE));
+
+ std::uniform_int_distribution<uint32_t> dist(kStartDate, kEndDate - 151);
+ uint32_t* o_orderdate = reinterpret_cast<uint32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) o_orderdate[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERPRIORITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERPRIORITY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERPRIORITY));
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_ORDERPRIORITY]);
+ std::uniform_int_distribution<int32_t> dist(0, kNumPriorities - 1);
+ char* o_orderpriority = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_ORDERPRIORITY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ const char* str = Priorities[dist(tld.rng)];
+ std::strncpy(o_orderpriority + i * byte_width, str, byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CLERK(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_CLERK].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CLERK));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_CLERK]);
+ int64_t max_clerk_id = static_cast<int64_t>(scale_factor_ * 1000);
+ std::uniform_int_distribution<int64_t> dist(1, max_clerk_id);
+ char* o_clerk = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_CLERK].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ const char* clerk = "Clerk#";
+ const size_t clerk_length = std::strlen(clerk);
+ int64_t clerk_number = dist(tld.rng);
+ char* output = o_clerk + i * byte_width;
+ std::strncpy(output, clerk, byte_width);
+ AppendNumberPaddedToNineDigits(output + clerk_length, clerk_number);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_SHIPPRIORITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_SHIPPRIORITY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_SHIPPRIORITY));
+ int32_t* o_shippriority = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_SHIPPRIORITY].array()->buffers[1]->mutable_data());
+ std::memset(o_shippriority, 0, tld.orders_to_generate * sizeof(int32_t));
+ }
+ return Status::OK();
+ }
+
+ Status O_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.orders[ORDERS::O_COMMENT],
+ g_text.GenerateComments(tld.orders_to_generate, 19, 78, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ Status GenerateRowCounts(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ std::uniform_int_distribution<int> length_dist(1, 7);
+ tld.lineitem_to_generate = 0;
+ tld.items_per_order.clear();
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ int length = length_dist(tld.rng);
+ tld.items_per_order.push_back(length);
+ tld.lineitem_to_generate += length;
+ }
+ int64_t num_batches =
+ (tld.first_batch_offset + tld.lineitem_to_generate + batch_size_ - 1) /
+ batch_size_;
+ tld.lineitem.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.lineitem) {
+ batch.resize(LINEITEM::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocateLineItemBufferIfNeeded(size_t thread_index, size_t ibatch, int column,
+ size_t& out_batch_offset) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.lineitem[ibatch][column].kind() == Datum::NONE) {
+ int32_t byte_width = arrow::internal::GetByteWidth(*kLineitemTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kLineitemTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.lineitem[ibatch][column] = std::move(ad);
+ out_batch_offset = 0;
+ }
+ if (ibatch == 0) out_batch_offset = tld.first_batch_offset;
+
+ return Status::OK();
+ }
+
+ Status L_ORDERKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_ORDERKEY]) {
+ tld.generated_lineitem[LINEITEM::L_ORDERKEY] = true;
+ RETURN_NOT_OK(O_ORDERKEY(thread_index));
+ const int32_t* o_orderkey = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_ORDERKEY, batch_offset));
+ int32_t* l_linenumber =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_ORDERKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_linenumber[batch_offset++] = o_orderkey[iorder];
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_ORDERKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_PARTKEY]) {
+ tld.generated_lineitem[LINEITEM::L_PARTKEY] = true;
+
+ size_t ibatch = 0;
+ int32_t max_partkey = static_cast<int32_t>(scale_factor_ * 200000);
+ std::uniform_int_distribution<int32_t> dist(1, max_partkey);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_PARTKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_partkey =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_partkey[batch_offset] = dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SUPPKEY]) {
+ tld.generated_lineitem[LINEITEM::L_SUPPKEY] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 3);
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset = 0;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SUPPKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_suppkey =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const int32_t* l_partkey = reinterpret_cast<const int32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ for (int64_t i = 0; i < next_run; i++) {
+ int32_t supplier = dist(tld.rng);
+ int32_t partkey = l_partkey[batch_offset];
+ // Fun fact: the parentheses for this expression are unbalanced in the TPC-H
+ // spec.
+ l_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SUPPKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINENUMBER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_LINENUMBER]) {
+ tld.generated_lineitem[LINEITEM::L_LINENUMBER] = true;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_LINENUMBER, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_linenumber =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_LINENUMBER]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++) {
+ l_linenumber[batch_offset++] = (iline + 1);
+ ARROW_DCHECK(1 <= (iline + 1) && (iline + 1) <= 7);
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINENUMBER].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_QUANTITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_QUANTITY]) {
+ tld.generated_lineitem[LINEITEM::L_QUANTITY] = true;
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int64_t> dist(1, 50);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_QUANTITY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_quantity =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_QUANTITY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t i = 0; i < next_run; i++) {
+ // Multiply by 100 because the type is decimal(12, 2), so the decimal goes after
+ // two digits
+ int64_t quantity = dist(tld.rng) * 100;
+ l_quantity[batch_offset++] = {quantity};
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_EXTENDEDPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE]) {
+ tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+ RETURN_NOT_OK(L_QUANTITY(thread_index));
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_EXTENDEDPRICE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ const int32_t* l_partkey = reinterpret_cast<const int32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ const Decimal128* l_quantity = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->buffers[1]->data());
+ Decimal128* l_extendedprice =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ int64_t partkey = static_cast<int64_t>(l_partkey[batch_offset]);
+ // Divide by 100 to recover the integer representation (not Decimal).
+ int64_t quantity = static_cast<int64_t>(l_quantity[batch_offset]) / 100;
+
+ // Spec says to divide by 100, but that happens automatically due to this being
+ // stored to two decimal points.
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ int64_t extended_price = retail_price * quantity;
+ l_extendedprice[batch_offset] = {extended_price};
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_DISCOUNT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_DISCOUNT]) {
+ tld.generated_lineitem[LINEITEM::L_DISCOUNT] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 10);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_DISCOUNT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_discount =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_DISCOUNT]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_discount[batch_offset] = {dist(tld.rng)};
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_TAX(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_TAX]) {
+ tld.generated_lineitem[LINEITEM::L_TAX] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 8);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_TAX, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_tax = reinterpret_cast<Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_tax[batch_offset] = {dist(tld.rng)};
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RETURNFLAG(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_RETURNFLAG]) {
+ tld.generated_lineitem[LINEITEM::L_RETURNFLAG] = true;
+ RETURN_NOT_OK(L_RECEIPTDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<uint32_t> dist;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_RETURNFLAG, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_returnflag =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_receiptdate = reinterpret_cast<const uint32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ if (l_receiptdate[batch_offset] <= kCurrentDate) {
+ uint32_t r = dist(tld.rng);
+ l_returnflag[batch_offset] = (r % 2 == 1) ? 'R' : 'A';
+ } else {
+ l_returnflag[batch_offset] = 'N';
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINESTATUS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_LINESTATUS]) {
+ tld.generated_lineitem[LINEITEM::L_LINESTATUS] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_LINESTATUS, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_linestatus =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_LINESTATUS]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_shipdate =
+ reinterpret_cast<const uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ if (l_shipdate[batch_offset] > kCurrentDate)
+ l_linestatus[batch_offset] = 'O';
+ else
+ l_linestatus[batch_offset] = 'F';
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPDATE]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPDATE] = true;
+ RETURN_NOT_OK(O_ORDERDATE(thread_index));
+ const int32_t* o_orderdate = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(1, 121);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SHIPDATE, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_shipdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_shipdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMITDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_COMMITDATE]) {
+ tld.generated_lineitem[LINEITEM::L_COMMITDATE] = true;
+ const int32_t* o_orderdate = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(30, 90);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_COMMITDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_commitdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_COMMITDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_commitdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMITDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RECEIPTDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_RECEIPTDATE]) {
+ tld.generated_lineitem[LINEITEM::L_RECEIPTDATE] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(1, 30);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_RECEIPTDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_receiptdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_shipdate =
+ reinterpret_cast<const uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_receiptdate[batch_offset] = l_shipdate[batch_offset] + dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPINSTRUCT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT] = true;
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPINSTRUCT]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumInstructions - 1);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_SHIPINSTRUCT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_shipinstruct =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ const char* str = Instructions[dist(tld.rng)];
+ // Note that we don't have to memset the buffer to 0 because strncpy pads each
+ // string with 0's anyway
+ std::strncpy(l_shipinstruct + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPMODE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPMODE]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPMODE] = true;
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPMODE]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumModes - 1);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SHIPMODE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_shipmode =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_SHIPMODE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ const char* str = Modes[dist(tld.rng)];
+ std::strncpy(l_shipmode + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPMODE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_COMMENT]) {
+ tld.generated_lineitem[LINEITEM::L_COMMENT] = true;
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ // Comments are kind of sneaky: we always generate the full batch and then just
+ // bump the length
+ if (tld.lineitem[ibatch][LINEITEM::L_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.lineitem[ibatch][LINEITEM::L_COMMENT],
+ g_text.GenerateComments(batch_size_, 10, 43, tld.rng));
+ batch_offset = 0;
+ }
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ batch_offset += next_run;
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMENT].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ std::vector<Datum> orders;
+ int64_t orders_to_generate;
+ int64_t orderkey_start;
+
+ std::vector<std::vector<Datum>> lineitem;
+ std::vector<int> items_per_order;
+ int64_t lineitem_to_generate;
+ int64_t first_batch_offset;
+ std::bitset<LINEITEM::kNumCols> generated_lineitem;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex orders_output_queue_mutex_;
+ std::mutex lineitem_output_queue_mutex_;
+ std::queue<ExecBatch> orders_output_queue_;
+ std::queue<ExecBatch> lineitem_output_queue_;
+ int64_t batch_size_;
+ float scale_factor_;
+ int64_t orders_rows_to_generate_;
+ int64_t orders_rows_generated_;
+ std::vector<int> orders_cols_;
+ std::vector<int> lineitem_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<size_t> orders_batches_generated_ = {0};
+ std::atomic<size_t> lineitem_batches_generated_ = {0};
+};
+
+class SupplierGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 10000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, gen_list_));
+
+ random::pcg32_fast rng;
+ std::uniform_int_distribution<int64_t> dist(0, rows_to_generate_ - 1);
+ size_t num_special_rows = static_cast<size_t>(5 * scale_factor_);
+ std::unordered_set<int64_t> good_rows_set;
+ while (good_rows_set.size() < num_special_rows) {
+ int64_t row = dist(rng);
+ good_rows_set.insert(row);
+ }
+ std::unordered_set<int64_t> bad_rows_set;
+ while (bad_rows_set.size() < num_special_rows) {
+ int64_t bad_row;
+ do {
+ bad_row = dist(rng);
+ } while (good_rows_set.find(bad_row) != good_rows_set.end());
+ bad_rows_set.insert(bad_row);
+ }
+ good_rows_.clear();
+ bad_rows_.clear();
+ good_rows_.insert(good_rows_.end(), good_rows_set.begin(), good_rows_set.end());
+ bad_rows_.insert(bad_rows_.end(), bad_rows_set.begin(), bad_rows_set.end());
+ std::sort(good_rows_.begin(), good_rows_.end());
+ std::sort(bad_rows_.begin(), bad_rows_.end());
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) tld.rng.seed(seq);
+
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(S_SUPPKEY) \
+ F(S_NAME) \
+ F(S_ADDRESS) \
+ F(S_NATIONKEY) \
+ F(S_PHONE) \
+ F(S_ACCTBAL) \
+ F(S_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct SUPPLIER {
+ enum {
+ FOR_EACH_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) {#col, SUPPLIER::col},
+ const std::unordered_map<std::string, int> kNameMap = {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kGenerators = {FOR_EACH_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(), fixed_size_binary(25), utf8(),
+ int32(), fixed_size_binary(15), decimal(12, 2),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index) {
+ if (done_.load()) return Status::OK();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.suppkey_start = rows_generated_.fetch_add(batch_size_);
+ if (tld.suppkey_start >= rows_to_generate_) return Status::OK();
+
+ tld.to_generate = std::min(batch_size_, rows_to_generate_ - tld.suppkey_start);
+
+ tld.batch.resize(SUPPLIER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for (int col : gen_list_) RETURN_NOT_OK(kGenerators[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for (size_t i = 0; i < gen_list_.size(); i++) {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_outputted_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_outputted_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if (is_last_batch) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ return Status::OK();
+ }
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(kTypes[column], tld.to_generate, {nullptr, std::move(buff)});
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status S_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_SUPPKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_SUPPKEY));
+ int32_t* s_suppkey = reinterpret_cast<int32_t*>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ s_suppkey[irow] = static_cast<int32_t>(tld.suppkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_NAME].kind() == Datum::NONE) {
+ RETURN_NOT_OK(S_SUPPKEY(thread_index));
+ const int32_t* s_suppkey = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->data());
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NAME));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[SUPPLIER::S_NAME]);
+ char* s_name = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_NAME].array()->buffers[1]->mutable_data());
+ // Look man, I'm just following the spec ok? Section 4.2.3 as of March 1 2022
+ const char* supplier = "Supplie#r";
+ const size_t supplier_length = std::strlen(supplier);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = s_name + byte_width * irow;
+ std::strncpy(out, supplier, byte_width);
+ AppendNumberPaddedToNineDigits(out + supplier_length, s_suppkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ADDRESS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_ADDRESS].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[SUPPLIER::S_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status S_NATIONKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_NATIONKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t* s_nationkey = reinterpret_cast<int32_t*>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status S_PHONE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_PHONE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(S_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[SUPPLIER::S_PHONE]);
+ const int32_t* s_nationkey = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->data());
+ char* s_phone = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_PHONE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ GeneratePhoneNumber(s_phone + irow * byte_width, tld.rng, s_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ACCTBAL(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_ACCTBAL].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_ACCTBAL));
+ Decimal128* s_acctbal = reinterpret_cast<Decimal128*>(
+ tld.batch[SUPPLIER::S_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_acctbal[irow] = {dist(tld.rng)};
+ }
+ return Status::OK();
+ }
+
+ Status S_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[SUPPLIER::S_COMMENT],
+ g_text.GenerateComments(tld.to_generate, 25, 100, tld.rng));
+ ModifyComments(thread_index, "Recommends", good_rows_);
+ ModifyComments(thread_index, "Complaints", bad_rows_);
+ }
+ return Status::OK();
+ }
+
+ void ModifyComments(size_t thread_index, const char* review,
+ const std::vector<int64_t>& indices) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ const int32_t* offsets = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[1]->data());
+ char* str = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[2]->mutable_data());
+ const char* customer = "Customer";
+ const int32_t customer_length = static_cast<int32_t>(std::strlen(customer));
+ const int32_t review_length = static_cast<int32_t>(std::strlen(review));
+
+ auto it = std::lower_bound(indices.begin(), indices.end(), tld.suppkey_start);
+ for (; it != indices.end() && *it < tld.suppkey_start + tld.to_generate; it++) {
+ int64_t idx_in_batch = *it - tld.suppkey_start;
+ char* out = str + offsets[idx_in_batch];
+ int32_t str_length = offsets[idx_in_batch + 1] - offsets[idx_in_batch];
+ std::uniform_int_distribution<int32_t> gap_dist(
+ 0, str_length - customer_length - review_length);
+ int32_t gap = gap_dist(tld.rng);
+ int32_t total_length = customer_length + gap + review_length;
+ std::uniform_int_distribution<int32_t> start_dist(0, str_length - total_length);
+ int32_t start = start_dist(tld.rng);
+ std::memcpy(out + start, customer, customer_length);
+ std::memcpy(out + start + customer_length + gap, review, review_length);
+ }
+ }
+
+ struct ThreadLocalData {
+ random::pcg32_fast rng;
+ int64_t suppkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+ std::vector<int64_t> good_rows_;
+ std::vector<int64_t> bad_rows_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_;
+ std::atomic<int64_t> rows_generated_;
+ float scale_factor_;
+ int64_t batch_size_;
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class PartGenerator : public TpchTableGenerator {
+ public:
+ explicit PartGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetPartOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch, gen_->NextPartBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->part_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class PartSuppGenerator : public TpchTableGenerator {
+ public:
+ explicit PartSuppGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetPartSuppOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextPartSuppBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->partsupp_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class CustomerGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, gen_list_));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) tld.rng.seed(seq);
+
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(C_CUSTKEY) \
+ F(C_NAME) \
+ F(C_ADDRESS) \
+ F(C_NATIONKEY) \
+ F(C_PHONE) \
+ F(C_ACCTBAL) \
+ F(C_MKTSEGMENT) \
+ F(C_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct CUSTOMER {
+ enum {
+ FOR_EACH_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) {#col, CUSTOMER::col},
+ const std::unordered_map<std::string, int> kNameMap = {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kGenerators = {FOR_EACH_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(),
+ utf8(),
+ utf8(),
+ int32(),
+ fixed_size_binary(15),
+ decimal(12, 2),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index) {
+ if (done_.load()) return Status::OK();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.custkey_start = rows_generated_.fetch_add(batch_size_);
+ if (tld.custkey_start >= rows_to_generate_) return Status::OK();
+
+ tld.to_generate = std::min(batch_size_, rows_to_generate_ - tld.custkey_start);
+
+ tld.batch.resize(CUSTOMER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for (int col : gen_list_) RETURN_NOT_OK(kGenerators[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for (size_t i = 0; i < gen_list_.size(); i++) {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_generated_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_generated_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if (is_last_batch) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ return Status::OK();
+ }
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(kTypes[column], tld.to_generate, {nullptr, std::move(buff)});
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status C_CUSTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_CUSTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_CUSTKEY));
+ int32_t* c_custkey = reinterpret_cast<int32_t*>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ c_custkey[irow] = static_cast<int32_t>(tld.custkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_NAME].kind() == Datum::NONE) {
+ RETURN_NOT_OK(C_CUSTKEY(thread_index));
+ const int32_t* c_custkey = reinterpret_cast<const int32_t*>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->data());
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ const char* customer = "Customer#";
+ const size_t customer_length = std::strlen(customer);
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ int num_digits = GetNumDigits(c_custkey[irow]);
+ int num_chars = std::max(num_digits, 9);
+ offsets[irow + 1] =
+ static_cast<int32_t>(offsets[irow] + num_chars + customer_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[tld.to_generate]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = str + offsets[irow];
+ std::memcpy(out, customer, customer_length);
+ AppendNumberPaddedToNineDigits(out + customer_length, c_custkey[irow]);
+ }
+ ArrayData ad(utf8(), tld.to_generate,
+ {nullptr, std::move(offset_buff), std::move(str_buff)});
+ tld.batch[CUSTOMER::C_NAME] = std::move(ad);
+ }
+ return Status::OK();
+ }
+
+ Status C_ADDRESS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_ADDRESS].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[CUSTOMER::C_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status C_NATIONKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_NATIONKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t* c_nationkey = reinterpret_cast<int32_t*>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status C_PHONE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_PHONE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(C_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[CUSTOMER::C_PHONE]);
+ const int32_t* c_nationkey = reinterpret_cast<const int32_t*>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->data());
+ char* c_phone = reinterpret_cast<char*>(
+ tld.batch[CUSTOMER::C_PHONE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ GeneratePhoneNumber(c_phone + irow * byte_width, tld.rng, c_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_ACCTBAL(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_ACCTBAL].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_ACCTBAL));
+ Decimal128* c_acctbal = reinterpret_cast<Decimal128*>(
+ tld.batch[CUSTOMER::C_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_acctbal[irow] = {dist(tld.rng)};
+ }
+ return Status::OK();
+ }
+
+ Status C_MKTSEGMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_MKTSEGMENT].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_MKTSEGMENT));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[CUSTOMER::C_MKTSEGMENT]);
+ char* c_mktsegment = reinterpret_cast<char*>(
+ tld.batch[CUSTOMER::C_MKTSEGMENT].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int32_t> dist(0, kNumSegments - 1);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = c_mktsegment + irow * byte_width;
+ int str_idx = dist(tld.rng);
+ std::strncpy(out, Segments[str_idx], byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[CUSTOMER::C_COMMENT],
+ g_text.GenerateComments(tld.to_generate, 29, 116, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ random::pcg32_fast rng;
+ int64_t custkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_{0};
+ std::atomic<int64_t> rows_generated_ = {0};
+ float scale_factor_{0};
+ int64_t batch_size_{0};
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class OrdersGenerator : public TpchTableGenerator {
+ public:
+ explicit OrdersGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetOrdersOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch, gen_->NextOrdersBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->orders_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class LineitemGenerator : public TpchTableGenerator {
+ public:
+ explicit LineitemGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetLineItemOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextLineItemBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->lineitem_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class NationGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, float /*scale_factor*/,
+ int64_t /*batch_size*/) override {
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, column_indices_));
+ rng_.seed(arrow_vendored::pcg_extras::seed_seq_from<std::random_device>{});
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t /*num_threads*/, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback /*schedule_task_callback*/) override {
+ std::shared_ptr<Buffer> N_NATIONKEY_buffer =
+ Buffer::Wrap(kNationKey, sizeof(kNationKey));
+ ArrayData N_NATIONKEY_arraydata(int32(), kRowCount,
+ {nullptr, std::move(N_NATIONKEY_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> N_NAME_buffer,
+ AllocateBuffer(kRowCount * kNameByteWidth));
+ char* N_NAME = reinterpret_cast<char*>(N_NAME_buffer->mutable_data());
+ for (size_t i = 0; i < kRowCount; i++)
+ std::strncpy(N_NAME + kNameByteWidth * i, kCountryNames[i], kNameByteWidth);
+ ArrayData N_NAME_arraydata(fixed_size_binary(kNameByteWidth), kRowCount,
+ {nullptr, std::move(N_NAME_buffer)});
+
+ std::shared_ptr<Buffer> N_REGIONKEY_buffer =
+ Buffer::Wrap(kRegionKey, sizeof(kRegionKey));
+ ArrayData N_REGIONKEY_arraydata(int32(), kRowCount,
+ {nullptr, std::move(N_REGIONKEY_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(Datum N_COMMENT_datum,
+ g_text.GenerateComments(kRowCount, 31, 114, rng_));
+
+ std::vector<Datum> fields = {
+ std::move(N_NATIONKEY_arraydata), std::move(N_NAME_arraydata),
+ std::move(N_REGIONKEY_arraydata), std::move(N_COMMENT_datum)};
+
+ std::vector<Datum> result;
+ for (const int& col : column_indices_) result.push_back(fields[col]);
+ ARROW_ASSIGN_OR_RAISE(ExecBatch batch, ExecBatch::Make(std::move(result)));
+ output_callback(std::move(batch));
+ finished_callback(static_cast<int64_t>(1));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ random::pcg32_fast rng_;
+
+ static constexpr size_t kRowCount = 25;
+ static constexpr int32_t kNameByteWidth = 25;
+ const int32_t kNationKey[kRowCount] = {0, 1, 2, 3, 4, 5, 6, 7, 8,
+ 9, 10, 11, 12, 13, 14, 15, 16, 17,
+ 18, 19, 20, 21, 22, 23, 24};
+ const char* kCountryNames[kRowCount] = {
+ "ALGERIA", "ARGENTINA", "BRAZIL", "CANADA", "EGYPT",
+ "ETHIOPIA", "FRANCE", "GERMANY", "INDIA", "INDONESIA",
+ "IRAN", "IRAQ", "JAPAN", "JORDAN", "KENYA",
+ "MOROCCO", "MOZAMBIQUE", "PERU", "CHINA", "ROMANIA",
+ "SAUDI ARABIA", "VIETNAM", "RUSSIA", "UNITED KINGDOM", "UNITED STATES"};
+ const int32_t kRegionKey[kRowCount] = {0, 1, 1, 1, 4, 0, 3, 3, 2, 2, 4, 4, 2,
+ 4, 0, 0, 0, 1, 2, 3, 4, 2, 3, 3, 1};
+
+ struct NATION {
+ enum {
Review comment:
no, i want NATION::NATIONKEY et. al to be implicitly convertible to int, so that I can use it as a name of indices. with enum class you have to add `static_cast` everywhere.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r836879877
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3483 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/formatting.h"
+#include "arrow/util/future.h"
+#include "arrow/util/io_util.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+using internal::GetRandomSeed;
+
+namespace compute {
+namespace internal {
+/*
+Architecture of the generator:
+ This is a multithreaded implementation of TPC-H's DBGen data generator. For each table
+ that doesn't depend on any other tables, it gets its own generator. For tables that
+ are dependent on each other (namely [ORDERS, LINEITEM] and [PART, PARTSUPP]), we
+implement a "double generator". Each generator is given a list of columns to generate.
+Columns are then generated lazily a batch at a time. If column A depends on another column
+B, column A will ensure that column B is always generated. We don't have to worry about
+breaking cycles of dependencies because the dependency graph is acyclic (so there's no
+code to even bother breaking cycles). Double generators work by maintaining two output
+queues. Batches for each of the tables are generated in sync, and batches that are not
+being output immediately are appended to the queue.
+
+ To generate a batch, we grab a lock, increment a counter, check the output queue (if
+applicable), and then call the necessary generator functions. To generate a column, we
+first check if it's already been generated. If not, we allocate the batch and then fill it
+according to the spec.
+
+ There are a few types of columns that get generated:
+ - Primary Keys: incrementing counters from 1 to N (with N being the number of rows).
+These are generated by incrementing a counter (this counter is gated under a lock).
+ - V-String: Random-length string of alphanumerics
+ - Phone Number: 2, 3, 3, and 4, digit numbers separated by -'s
+ - Random Numbers: random numbers within some range
+ - Expressions: expressions based on some other columns
+ Please consult the spec for intended behavior of each individual column. Columns are
+generated by a function of the same name (so e.g. the column PS_PARTKEY is generated by a
+function PS_PARTKEY).
+*/
+
+namespace {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+std::uniform_int_distribution<int64_t> kSeedDist(std::numeric_limits<int64_t>::min(),
+ std::numeric_limits<int64_t>::max());
+} // namespace
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, double scale_factor,
+ int64_t batch_size, int64_t seed) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ int64_t seed_ = {0};
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ Unreachable("Positive 64-bit integer should never have more than 18 digits");
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ size_t kPad = 9;
+ out += std::max(kPad, static_cast<size_t>(GetNumDigits(x)));
+ arrow::internal::detail::FormatAllDigitsLeftPadded(x, kPad, '0', &out);
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+
+ out += 15;
+ arrow::internal::detail::FormatAllDigits(l3, &out);
+ *(--out) = '-';
+ arrow::internal::detail::FormatAllDigits(l2, &out);
+ *(--out) = '-';
+ arrow::internal::detail::FormatAllDigits(l1, &out);
+ *(--out) = '-';
+ arrow::internal::detail::FormatTwoDigits(country_code, &out);
+}
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, double scale_factor, int64_t seed) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ thread_local_data_.resize(num_threads);
+ random::pcg64_fast seed_rng(seed);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(kSeedDist(seed_rng));
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ if (!partsupp_output_queue_.empty()) {
+ ExecBatch result = std::move(partsupp_output_queue_.front());
+ partsupp_output_queue_.pop();
+ return std::move(result);
+ }
+ }
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+ if (!part_cols_.empty()) {
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch part_batch,
+ ExecBatch::Make(std::move(part_result)));
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ part_output_queue_.emplace(std::move(part_batch));
+ }
+ }
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ for (size_t i = 1; i < partsupp_results.size(); i++)
+ partsupp_output_queue_.emplace(std::move(partsupp_results[i]));
+ }
+ return std::move(partsupp_results[0]);
+ }
+
+ private:
+#define FOR_EACH_PART_COLUMN(F) \
+ F(P_PARTKEY) \
+ F(P_NAME) \
+ F(P_MFGR) \
+ F(P_BRAND) \
+ F(P_TYPE) \
+ F(P_SIZE) \
+ F(P_CONTAINER) \
+ F(P_RETAILPRICE) \
+ F(P_COMMENT)
+
+#define FOR_EACH_PARTSUPP_COLUMN(F) \
+ F(PS_PARTKEY) \
+ F(PS_SUPPKEY) \
+ F(PS_AVAILQTY) \
+ F(PS_SUPPLYCOST) \
+ F(PS_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct PART {
+ enum {
+ FOR_EACH_PART_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct PARTSUPP {
+ enum {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, PART::col},
+ const std::unordered_map<std::string, int> kPartNameMap = {
+ FOR_EACH_PART_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, PARTSUPP::col},
+ const std::unordered_map<std::string, int> kPartsuppNameMap = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kPartGenerators = {FOR_EACH_PART_COLUMN(MAKE_FN_ARRAY)};
+ std::vector<GenerateColumnFn> kPartsuppGenerators = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kPartTypes = {
+ int32(), utf8(), fixed_size_binary(25), fixed_size_binary(10),
+ utf8(), int32(), fixed_size_binary(10), decimal(12, 2),
+ utf8(),
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kPartsuppTypes = {
+ int32(), int32(), int32(), decimal(12, 2), utf8(),
+ };
+
+ Status AllocatePartBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.part[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.part_to_generate * byte_width));
+ ArrayData ad(kPartTypes[column], tld.part_to_generate, {nullptr, std::move(buff)});
+ tld.part[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status P_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_PARTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_PARTKEY));
+ int32_t* p_partkey = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) {
+ p_partkey[i] = static_cast<int32_t>(tld.partkey_start + i + 1);
+ ARROW_DCHECK(1 <= p_partkey[i] && p_partkey[i] <= part_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_NAME].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(
+ 0, static_cast<uint8_t>(kNumNameParts - 1));
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dist(tld.rng));
+ tld.string_indices[irow * 5 + ipart] = name_part_index;
+ string_length += std::strlen(NameParts[name_part_index]);
+ }
+ // Add 4 because there is a space between each word (i.e. four spaces)
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length + 4);
+ }
+ // Add an extra byte for the space after in the very last string.
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate] + 1));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 5 + ipart];
+ const char* part = NameParts[name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ *row++ = ' ';
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_NAME], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_NAME] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_MFGR(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_MFGR].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* manufacturer = "Manufacturer#";
+ const size_t manufacturer_length = std::strlen(manufacturer);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_MFGR));
+ char* p_mfgr = reinterpret_cast<char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ std::strncpy(p_mfgr + byte_width * irow, manufacturer, byte_width);
+ char mfgr_id = '0' + dist(tld.rng);
+ *(p_mfgr + byte_width * irow + manufacturer_length) = mfgr_id;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_BRAND(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_BRAND].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_MFGR(thread_index));
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* brand = "Brand#";
+ const size_t brand_length = std::strlen(brand);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_BRAND));
+ const char* p_mfgr = reinterpret_cast<const char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->data());
+ char* p_brand = reinterpret_cast<char*>(
+ tld.part[PART::P_BRAND].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_BRAND]);
+ int32_t mfgr_byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ const size_t mfgr_id_offset = std::strlen("Manufacturer#");
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = p_brand + byte_width * irow;
+ char mfgr_id = *(p_mfgr + irow * mfgr_byte_width + mfgr_id_offset);
+ char brand_id = '0' + dist(tld.rng);
+ std::strncpy(row, brand, byte_width);
+ *(row + brand_length) = mfgr_id;
+ *(row + brand_length + 1) = brand_id;
+ irow += 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_TYPE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_TYPE].kind() == Datum::NONE) {
+ using D = std::uniform_int_distribution<uint32_t>;
+ D dists[] = {
+ D{0, static_cast<uint8_t>(kNumTypes_1 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_2 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_3 - 1)},
+ };
+
+ const char** types[] = {Types_1, Types_2, Types_3};
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dists[ipart](tld.rng));
+ tld.string_indices[irow * 3 + ipart] = name_part_index;
+ string_length += std::strlen(types[ipart][name_part_index]);
+ }
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate]));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 3 + ipart];
+ const char* part = types[ipart][name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_TYPE], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_TYPE] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_SIZE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_SIZE].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(1, 50);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_SIZE));
+ int32_t* p_size = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_SIZE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) p_size[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status P_CONTAINER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_CONTAINER].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist1(
+ 0, static_cast<uint8_t>(kNumContainers_1 - 1));
+ std::uniform_int_distribution<int> dist2(
+ 0, static_cast<uint8_t>(kNumContainers_2 - 1));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_CONTAINER));
+ char* p_container = reinterpret_cast<char*>(
+ tld.part[PART::P_CONTAINER].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_CONTAINER]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int container1_idx = dist1(tld.rng);
+ int container2_idx = dist2(tld.rng);
+ const char* container1 = Containers_1[container1_idx];
+ const char* container2 = Containers_2[container2_idx];
+ size_t container1_length = std::strlen(container1);
+ size_t container2_length = std::strlen(container2);
+
+ char* row = p_container + byte_width * irow;
+ std::strncpy(row, container1, byte_width);
+ std::memcpy(row + container1_length, container2, container2_length);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_RETAILPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_RETAILPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_RETAILPRICE));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+ Decimal128* p_retailprice = reinterpret_cast<Decimal128*>(
+ tld.part[PART::P_RETAILPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int32_t partkey = p_partkey[irow];
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ p_retailprice[irow] = {retail_price};
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.part[PART::P_COMMENT],
+ g_text.GenerateComments(tld.part_to_generate, 5, 22, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ int64_t PartsuppBatchesToGenerate(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t num_batches = (ps_to_generate + batch_size_ - 1) / batch_size_;
+ return num_batches;
+ }
+
+ Status InitPartsupp(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.generated_partsupp.reset();
+ int64_t num_batches = PartsuppBatchesToGenerate(thread_index);
+ tld.partsupp.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.partsupp) {
+ batch.resize(PARTSUPP::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocatePartSuppBatch(size_t thread_index, size_t ibatch, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartsuppTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kPartsuppTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.partsupp[ibatch][column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status PS_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_PARTKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_PARTKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_PARTKEY));
+ int32_t* ps_partkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run; ipartsupp++, irun++)
+ ps_partkey[batch_offset++] = p_partkey[ipart];
+
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPKEY));
+ int32_t* ps_suppkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run;
+ ipartsupp++, irun++) {
+ int32_t supplier = static_cast<int32_t>(ipartsupp);
+ int32_t partkey = p_partkey[ipart];
+ // Magic formula from TPCH spec.
+ ps_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_AVAILQTY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_AVAILQTY]) {
+ tld.generated_partsupp[PARTSUPP::PS_AVAILQTY] = true;
+ std::uniform_int_distribution<int32_t> dist(1, 9999);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_AVAILQTY));
+ int32_t* ps_availqty =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for (int64_t irun = 0; irun < next_run; irun++) ps_availqty[irun] = dist(tld.rng);
+
+ tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPLYCOST(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST] = true;
+ std::uniform_int_distribution<int64_t> dist(100, 100000);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(
+ AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPLYCOST));
+ Decimal128* ps_supplycost =
+ reinterpret_cast<Decimal128*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for (int64_t irun = 0; irun < next_run; irun++)
+ ps_supplycost[irun] = {dist(tld.rng)};
+
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PARTSUPP::PS_COMMENT].kind() == Datum::NONE) {
+ int64_t irow = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ int64_t num_rows = std::min(batch_size_, ps_to_generate - irow);
+ ARROW_ASSIGN_OR_RAISE(tld.partsupp[ibatch][PARTSUPP::PS_COMMENT],
+ g_text.GenerateComments(num_rows, 49, 198, tld.rng));
+ irow += num_rows;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ std::vector<Datum> part;
+ std::vector<int8_t> string_indices;
+ int64_t part_to_generate{0};
+ int64_t partkey_start{0};
+
+ std::vector<std::vector<Datum>> partsupp;
+ std::bitset<PARTSUPP::kNumCols> generated_partsupp;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex part_output_queue_mutex_;
+ std::mutex partsupp_output_queue_mutex_;
+ std::queue<ExecBatch> part_output_queue_;
+ std::queue<ExecBatch> partsupp_output_queue_;
+ int64_t batch_size_{0};
+ double scale_factor_{0};
+ int64_t part_rows_to_generate_{0};
+ int64_t part_rows_generated_{0};
+ std::vector<int> part_cols_;
+ std::vector<int> partsupp_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<int64_t> part_batches_generated_ = {0};
+ std::atomic<int64_t> partsupp_batches_generated_ = {0};
+ static constexpr int64_t kPartSuppRowsPerPart = 4;
+};
+
+class OrdersAndLineItemGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, double scale_factor, int64_t seed) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ thread_local_data_.resize(num_threads);
+ random::pcg64_fast seed_rng(seed);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ tld.items_per_order.resize(batch_size_);
+ tld.rng.seed(kSeedDist(seed_rng));
+ }
+ orders_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000 * 10);
+ }
+ return Status::OK();
+ }
+
+ int64_t orders_batches_generated() const { return orders_batches_generated_.load(); }
+
+ int64_t lineitem_batches_generated() const {
+ return lineitem_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetOrdersOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kOrdersTypes, kOrdersNameMap, orders_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetLineItemOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kLineitemTypes, kLineitemNameMap, lineitem_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextOrdersBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if (!orders_output_queue_.empty()) {
+ ExecBatch batch = std::move(orders_output_queue_.front());
+ orders_output_queue_.pop();
+ return std::move(batch);
+ } else if (orders_rows_generated_ == orders_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate =
+ std::min(batch_size_, orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.first_batch_offset = 0;
+ tld.generated_lineitem.reset();
+
+ for (int col : orders_cols_) RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for (int col : lineitem_cols_) RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for (size_t i = 0; i < orders_cols_.size(); i++) {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ if (!lineitem_cols_.empty()) {
+ std::vector<ExecBatch> lineitem_results;
+ for (size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++) {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for (size_t icol = 0; icol < lineitem_cols_.size(); icol++) {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(lineitem_output_queue_mutex_);
+ for (ExecBatch& eb : lineitem_results) {
+ lineitem_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(orders_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextLineItemBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ExecBatch queued;
+ bool from_queue = false;
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ if (!lineitem_output_queue_.empty()) {
+ queued = std::move(lineitem_output_queue_.front());
+ lineitem_output_queue_.pop();
+ from_queue = true;
+ }
+ }
+ tld.first_batch_offset = 0;
+ if (from_queue) {
+ ARROW_DCHECK(queued.length <= batch_size_);
+ tld.first_batch_offset = queued.length;
+ if (queued.length == batch_size_) return std::move(queued);
+ }
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if (orders_rows_generated_ == orders_rows_to_generate_) {
+ if (from_queue) return std::move(queued);
+ return util::nullopt;
+ }
+
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate =
+ std::min(batch_size_, orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1ll);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.generated_lineitem.reset();
+ if (from_queue) {
+ lineitem_batches_generated_.fetch_sub(1);
+ for (size_t i = 0; i < lineitem_cols_.size(); i++)
+ if (tld.lineitem[0][lineitem_cols_[i]].kind() == Datum::NONE)
+ tld.lineitem[0][lineitem_cols_[i]] = std::move(queued[i]);
+ }
+
+ for (int col : orders_cols_) RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for (int col : lineitem_cols_) RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ if (!orders_cols_.empty()) {
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for (size_t i = 0; i < orders_cols_.size(); i++) {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch orders_batch,
+ ExecBatch::Make(std::move(orders_result)));
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ orders_output_queue_.emplace(std::move(orders_batch));
+ }
+ }
+ std::vector<ExecBatch> lineitem_results;
+ for (size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++) {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for (size_t icol = 0; icol < lineitem_cols_.size(); icol++) {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ lineitem_batches_generated_.fetch_add(static_cast<int64_t>(lineitem_results.size()));
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ for (size_t i = 1; i < lineitem_results.size(); i++)
+ lineitem_output_queue_.emplace(std::move(lineitem_results[i]));
+ }
+ return std::move(lineitem_results[0]);
+ }
+
+ private:
+#define FOR_EACH_ORDERS_COLUMN(F) \
+ F(O_ORDERKEY) \
+ F(O_CUSTKEY) \
+ F(O_ORDERSTATUS) \
+ F(O_TOTALPRICE) \
+ F(O_ORDERDATE) \
+ F(O_ORDERPRIORITY) \
+ F(O_CLERK) \
+ F(O_SHIPPRIORITY) \
+ F(O_COMMENT)
+
+#define FOR_EACH_LINEITEM_COLUMN(F) \
+ F(L_ORDERKEY) \
+ F(L_PARTKEY) \
+ F(L_SUPPKEY) \
+ F(L_LINENUMBER) \
+ F(L_QUANTITY) \
+ F(L_EXTENDEDPRICE) \
+ F(L_DISCOUNT) \
+ F(L_TAX) \
+ F(L_RETURNFLAG) \
+ F(L_LINESTATUS) \
+ F(L_SHIPDATE) \
+ F(L_COMMITDATE) \
+ F(L_RECEIPTDATE) \
+ F(L_SHIPINSTRUCT) \
+ F(L_SHIPMODE) \
+ F(L_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct ORDERS {
+ enum {
+ FOR_EACH_ORDERS_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct LINEITEM {
+ enum {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, ORDERS::col},
+ const std::unordered_map<std::string, int> kOrdersNameMap = {
+ FOR_EACH_ORDERS_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, LINEITEM::col},
+ const std::unordered_map<std::string, int> kLineitemNameMap = {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ const std::vector<GenerateColumnFn> kOrdersGenerators = {
+ FOR_EACH_ORDERS_COLUMN(MAKE_FN_ARRAY)};
+ const std::vector<GenerateColumnFn> kLineitemGenerators = {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kOrdersTypes = {int32(),
+ int32(),
+ fixed_size_binary(1),
+ decimal(12, 2),
+ date32(),
+ fixed_size_binary(15),
+ fixed_size_binary(15),
+ int32(),
+ utf8()};
+
+ const std::vector<std::shared_ptr<DataType>> kLineitemTypes = {
+ int32(),
+ int32(),
+ int32(),
+ int32(),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ fixed_size_binary(1),
+ fixed_size_binary(1),
+ date32(),
+ date32(),
+ date32(),
+ fixed_size_binary(25),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status AllocateOrdersBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.orders[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.orders_to_generate * byte_width));
+ ArrayData ad(kOrdersTypes[column], tld.orders_to_generate,
+ {nullptr, std::move(buff)});
+ tld.orders[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status O_ORDERKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERKEY));
+ int32_t* o_orderkey = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ int32_t orderkey_index = static_cast<int32_t>(tld.orderkey_start + i);
+ int32_t index_of_run = orderkey_index / 8;
+ int32_t index_in_run = orderkey_index % 8;
+ o_orderkey[i] = (index_of_run * 32 + index_in_run + 1);
+ ARROW_DCHECK(1 <= o_orderkey[i] && o_orderkey[i] <= 4 * orders_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CUSTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_CUSTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CUSTKEY));
+
+ // Spec says it must be a random number between 1 and SF*150000 that is not
+ // divisible by 3. Rather than repeatedly generating numbers until we get to
+ // a non-divisible-by-3 number, we just generate a number between
+ // 0 and SF * 50000 - 1, multiply by 3, and then add either 1 or 2.
+ int32_t sf_50k = static_cast<int32_t>(scale_factor_ * 50000);
+ std::uniform_int_distribution<int32_t> base_dist(0, sf_50k - 1);
+ std::uniform_int_distribution<int32_t> offset_dist(1, 2);
+ int32_t* o_custkey = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_CUSTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++)
+ o_custkey[i] = 3 * base_dist(tld.rng) + offset_dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERSTATUS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERSTATUS].kind() == Datum::NONE) {
+ RETURN_NOT_OK(L_LINESTATUS(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERSTATUS));
+
+ char* o_orderstatus = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_ORDERSTATUS].array()->buffers[1]->mutable_data());
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ bool all_f = true;
+ bool all_o = true;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ const char* l_linestatus = reinterpret_cast<const char*>(
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->buffers[1]->data());
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++, batch_offset++) {
+ all_f &= l_linestatus[batch_offset] == 'F';
+ all_o &= l_linestatus[batch_offset] == 'O';
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ ARROW_DCHECK(!(all_f && all_o));
+ if (all_f)
+ o_orderstatus[iorder] = 'F';
+ else if (all_o)
+ o_orderstatus[iorder] = 'O';
+ else
+ o_orderstatus[iorder] = 'P';
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_TOTALPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_TOTALPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(L_EXTENDEDPRICE(thread_index));
+ RETURN_NOT_OK(L_TAX(thread_index));
+ RETURN_NOT_OK(L_DISCOUNT(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_TOTALPRICE));
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ int64_t sum = 0;
+ Decimal128* o_totalprice = reinterpret_cast<Decimal128*>(
+ tld.orders[ORDERS::O_TOTALPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ const Decimal128* l_extendedprice = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->buffers[1]->data());
+ const Decimal128* l_tax = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->data());
+ const Decimal128* l_discount = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->buffers[1]->data());
+
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++, batch_offset++) {
+ int64_t eprice = static_cast<int64_t>(l_extendedprice[batch_offset]);
+ int64_t tax = static_cast<int64_t>(l_tax[batch_offset]);
+ int64_t discount = static_cast<int64_t>(l_discount[batch_offset]);
+ sum += (eprice * (100 + tax) * (100 - discount));
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ sum /= 100 * 100;
+ o_totalprice[iorder] = {sum};
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERDATE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERDATE));
+
+ std::uniform_int_distribution<uint32_t> dist(kStartDate, kEndDate - 151);
+ uint32_t* o_orderdate = reinterpret_cast<uint32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) o_orderdate[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERPRIORITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERPRIORITY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERPRIORITY));
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_ORDERPRIORITY]);
+ std::uniform_int_distribution<int32_t> dist(0, kNumPriorities - 1);
+ char* o_orderpriority = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_ORDERPRIORITY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ const char* str = Priorities[dist(tld.rng)];
+ std::strncpy(o_orderpriority + i * byte_width, str, byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CLERK(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_CLERK].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CLERK));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_CLERK]);
+ int64_t max_clerk_id = static_cast<int64_t>(scale_factor_ * 1000);
+ std::uniform_int_distribution<int64_t> dist(1, max_clerk_id);
+ char* o_clerk = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_CLERK].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ const char* clerk = "Clerk#";
+ const size_t clerk_length = std::strlen(clerk);
+ int64_t clerk_number = dist(tld.rng);
+ char* output = o_clerk + i * byte_width;
+ std::strncpy(output, clerk, byte_width);
+ AppendNumberPaddedToNineDigits(output + clerk_length, clerk_number);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_SHIPPRIORITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_SHIPPRIORITY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_SHIPPRIORITY));
+ int32_t* o_shippriority = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_SHIPPRIORITY].array()->buffers[1]->mutable_data());
+ std::memset(o_shippriority, 0, tld.orders_to_generate * sizeof(int32_t));
+ }
+ return Status::OK();
+ }
+
+ Status O_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.orders[ORDERS::O_COMMENT],
+ g_text.GenerateComments(tld.orders_to_generate, 19, 78, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ Status GenerateRowCounts(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ std::uniform_int_distribution<int> length_dist(1, 7);
+ tld.lineitem_to_generate = 0;
+ tld.items_per_order.clear();
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ int length = length_dist(tld.rng);
+ tld.items_per_order.push_back(length);
+ tld.lineitem_to_generate += length;
+ }
+ int64_t num_batches =
+ (tld.first_batch_offset + tld.lineitem_to_generate + batch_size_ - 1) /
+ batch_size_;
+ tld.lineitem.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.lineitem) {
+ batch.resize(LINEITEM::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocateLineItemBufferIfNeeded(size_t thread_index, size_t ibatch, int column,
+ size_t& out_batch_offset) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.lineitem[ibatch][column].kind() == Datum::NONE) {
+ int32_t byte_width = arrow::internal::GetByteWidth(*kLineitemTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kLineitemTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.lineitem[ibatch][column] = std::move(ad);
+ out_batch_offset = 0;
+ }
+ if (ibatch == 0) out_batch_offset = tld.first_batch_offset;
+
+ return Status::OK();
+ }
+
+ Status L_ORDERKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_ORDERKEY]) {
+ tld.generated_lineitem[LINEITEM::L_ORDERKEY] = true;
+ RETURN_NOT_OK(O_ORDERKEY(thread_index));
+ const int32_t* o_orderkey = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_ORDERKEY, batch_offset));
+ int32_t* l_linenumber =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_ORDERKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_linenumber[batch_offset++] = o_orderkey[iorder];
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_ORDERKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_PARTKEY]) {
+ tld.generated_lineitem[LINEITEM::L_PARTKEY] = true;
+
+ size_t ibatch = 0;
+ int32_t max_partkey = static_cast<int32_t>(scale_factor_ * 200000);
+ std::uniform_int_distribution<int32_t> dist(1, max_partkey);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_PARTKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_partkey =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_partkey[batch_offset] = dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SUPPKEY]) {
+ tld.generated_lineitem[LINEITEM::L_SUPPKEY] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 3);
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset = 0;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SUPPKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_suppkey =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const int32_t* l_partkey = reinterpret_cast<const int32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ for (int64_t i = 0; i < next_run; i++) {
+ int32_t supplier = dist(tld.rng);
+ int32_t partkey = l_partkey[batch_offset];
+ // Fun fact: the parentheses for this expression are unbalanced in the TPC-H
+ // spec.
+ l_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SUPPKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINENUMBER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_LINENUMBER]) {
+ tld.generated_lineitem[LINEITEM::L_LINENUMBER] = true;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_LINENUMBER, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_linenumber =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_LINENUMBER]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++) {
+ l_linenumber[batch_offset++] = (iline + 1);
+ ARROW_DCHECK(1 <= (iline + 1) && (iline + 1) <= 7);
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINENUMBER].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_QUANTITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_QUANTITY]) {
+ tld.generated_lineitem[LINEITEM::L_QUANTITY] = true;
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int64_t> dist(1, 50);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_QUANTITY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_quantity =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_QUANTITY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t i = 0; i < next_run; i++) {
+ // Multiply by 100 because the type is decimal(12, 2), so the decimal goes after
+ // two digits
+ int64_t quantity = dist(tld.rng) * 100;
+ l_quantity[batch_offset++] = {quantity};
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_EXTENDEDPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE]) {
+ tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+ RETURN_NOT_OK(L_QUANTITY(thread_index));
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_EXTENDEDPRICE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ const int32_t* l_partkey = reinterpret_cast<const int32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ const Decimal128* l_quantity = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->buffers[1]->data());
+ Decimal128* l_extendedprice =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ int64_t partkey = static_cast<int64_t>(l_partkey[batch_offset]);
+ // Divide by 100 to recover the integer representation (not Decimal).
+ int64_t quantity = static_cast<int64_t>(l_quantity[batch_offset]) / 100;
+
+ // Spec says to divide by 100, but that happens automatically due to this being
+ // stored to two decimal points.
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ int64_t extended_price = retail_price * quantity;
+ l_extendedprice[batch_offset] = {extended_price};
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_DISCOUNT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_DISCOUNT]) {
+ tld.generated_lineitem[LINEITEM::L_DISCOUNT] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 10);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_DISCOUNT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_discount =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_DISCOUNT]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_discount[batch_offset] = {dist(tld.rng)};
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_TAX(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_TAX]) {
+ tld.generated_lineitem[LINEITEM::L_TAX] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 8);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_TAX, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_tax = reinterpret_cast<Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_tax[batch_offset] = {dist(tld.rng)};
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RETURNFLAG(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_RETURNFLAG]) {
+ tld.generated_lineitem[LINEITEM::L_RETURNFLAG] = true;
+ RETURN_NOT_OK(L_RECEIPTDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<uint32_t> dist;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_RETURNFLAG, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_returnflag =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_receiptdate = reinterpret_cast<const uint32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ if (l_receiptdate[batch_offset] <= kCurrentDate) {
+ uint32_t r = dist(tld.rng);
+ l_returnflag[batch_offset] = (r % 2 == 1) ? 'R' : 'A';
+ } else {
+ l_returnflag[batch_offset] = 'N';
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINESTATUS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_LINESTATUS]) {
+ tld.generated_lineitem[LINEITEM::L_LINESTATUS] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_LINESTATUS, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_linestatus =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_LINESTATUS]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_shipdate =
+ reinterpret_cast<const uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ if (l_shipdate[batch_offset] > kCurrentDate)
+ l_linestatus[batch_offset] = 'O';
+ else
+ l_linestatus[batch_offset] = 'F';
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPDATE]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPDATE] = true;
+ RETURN_NOT_OK(O_ORDERDATE(thread_index));
+ const int32_t* o_orderdate = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(1, 121);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SHIPDATE, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_shipdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_shipdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMITDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_COMMITDATE]) {
+ tld.generated_lineitem[LINEITEM::L_COMMITDATE] = true;
+ const int32_t* o_orderdate = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(30, 90);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_COMMITDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_commitdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_COMMITDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_commitdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMITDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RECEIPTDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_RECEIPTDATE]) {
+ tld.generated_lineitem[LINEITEM::L_RECEIPTDATE] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(1, 30);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_RECEIPTDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_receiptdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_shipdate =
+ reinterpret_cast<const uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_receiptdate[batch_offset] = l_shipdate[batch_offset] + dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPINSTRUCT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT] = true;
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPINSTRUCT]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumInstructions - 1);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_SHIPINSTRUCT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_shipinstruct =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ const char* str = Instructions[dist(tld.rng)];
+ // Note that we don't have to memset the buffer to 0 because strncpy pads each
+ // string with 0's anyway
+ std::strncpy(l_shipinstruct + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPMODE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPMODE]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPMODE] = true;
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPMODE]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumModes - 1);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SHIPMODE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_shipmode =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_SHIPMODE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ const char* str = Modes[dist(tld.rng)];
+ std::strncpy(l_shipmode + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPMODE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_COMMENT]) {
+ tld.generated_lineitem[LINEITEM::L_COMMENT] = true;
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ // Comments are kind of sneaky: we always generate the full batch and then just
+ // bump the length
+ if (tld.lineitem[ibatch][LINEITEM::L_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.lineitem[ibatch][LINEITEM::L_COMMENT],
+ g_text.GenerateComments(batch_size_, 10, 43, tld.rng));
+ batch_offset = 0;
+ }
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ batch_offset += next_run;
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMENT].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ std::vector<Datum> orders;
+ int64_t orders_to_generate;
+ int64_t orderkey_start;
+
+ std::vector<std::vector<Datum>> lineitem;
+ std::vector<int> items_per_order;
+ int64_t lineitem_to_generate;
+ int64_t first_batch_offset;
+ std::bitset<LINEITEM::kNumCols> generated_lineitem;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex orders_output_queue_mutex_;
+ std::mutex lineitem_output_queue_mutex_;
+ std::queue<ExecBatch> orders_output_queue_;
+ std::queue<ExecBatch> lineitem_output_queue_;
+ int64_t batch_size_;
+ double scale_factor_;
+ int64_t orders_rows_to_generate_;
+ int64_t orders_rows_generated_;
+ std::vector<int> orders_cols_;
+ std::vector<int> lineitem_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<size_t> orders_batches_generated_ = {0};
+ std::atomic<size_t> lineitem_batches_generated_ = {0};
+};
+
+class SupplierGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, double scale_factor, int64_t batch_size,
+ int64_t seed) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 10000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, gen_list_));
+
+ random::pcg64_fast seed_rng(seed);
+ seed_ = kSeedDist(seed_rng);
+
+ random::pcg32_fast rng(kSeedDist(seed_rng));
+ std::uniform_int_distribution<int64_t> dist(0, rows_to_generate_ - 1);
+ size_t num_special_rows = static_cast<size_t>(5 * scale_factor_);
+ std::unordered_set<int64_t> good_rows_set;
+ while (good_rows_set.size() < num_special_rows) {
+ int64_t row = dist(rng);
+ good_rows_set.insert(row);
+ }
+ std::unordered_set<int64_t> bad_rows_set;
+ while (bad_rows_set.size() < num_special_rows) {
+ int64_t bad_row;
+ do {
+ bad_row = dist(rng);
+ } while (good_rows_set.find(bad_row) != good_rows_set.end());
+ bad_rows_set.insert(bad_row);
+ }
+ good_rows_.clear();
+ bad_rows_.clear();
+ good_rows_.insert(good_rows_.end(), good_rows_set.begin(), good_rows_set.end());
+ bad_rows_.insert(bad_rows_.end(), bad_rows_set.begin(), bad_rows_set.end());
+ std::sort(good_rows_.begin(), good_rows_.end());
+ std::sort(bad_rows_.begin(), bad_rows_.end());
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) tld.rng.seed(GetRandomSeed());
+
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(S_SUPPKEY) \
+ F(S_NAME) \
+ F(S_ADDRESS) \
+ F(S_NATIONKEY) \
+ F(S_PHONE) \
+ F(S_ACCTBAL) \
+ F(S_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct SUPPLIER {
+ enum {
+ FOR_EACH_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) {#col, SUPPLIER::col},
+ const std::unordered_map<std::string, int> kNameMap = {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kGenerators = {FOR_EACH_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(), fixed_size_binary(25), utf8(),
+ int32(), fixed_size_binary(15), decimal(12, 2),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index) {
+ if (done_.load()) return Status::OK();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.suppkey_start = rows_generated_.fetch_add(batch_size_);
+ if (tld.suppkey_start >= rows_to_generate_) return Status::OK();
+
+ tld.to_generate = std::min(batch_size_, rows_to_generate_ - tld.suppkey_start);
+
+ tld.batch.resize(SUPPLIER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for (int col : gen_list_) RETURN_NOT_OK(kGenerators[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for (size_t i = 0; i < gen_list_.size(); i++) {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_outputted_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_outputted_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if (is_last_batch) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ return Status::OK();
+ }
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(kTypes[column], tld.to_generate, {nullptr, std::move(buff)});
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status S_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_SUPPKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_SUPPKEY));
+ int32_t* s_suppkey = reinterpret_cast<int32_t*>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ s_suppkey[irow] = static_cast<int32_t>(tld.suppkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_NAME].kind() == Datum::NONE) {
+ RETURN_NOT_OK(S_SUPPKEY(thread_index));
+ const int32_t* s_suppkey = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->data());
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NAME));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[SUPPLIER::S_NAME]);
+ char* s_name = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_NAME].array()->buffers[1]->mutable_data());
+ // Look man, I'm just following the spec ok? Section 4.2.3 as of March 1 2022
+ const char* supplier = "Supplie#r";
+ const size_t supplier_length = std::strlen(supplier);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = s_name + byte_width * irow;
+ std::strncpy(out, supplier, byte_width);
+ AppendNumberPaddedToNineDigits(out + supplier_length, s_suppkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ADDRESS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_ADDRESS].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[SUPPLIER::S_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status S_NATIONKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_NATIONKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t* s_nationkey = reinterpret_cast<int32_t*>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status S_PHONE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_PHONE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(S_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[SUPPLIER::S_PHONE]);
+ const int32_t* s_nationkey = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->data());
+ char* s_phone = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_PHONE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ GeneratePhoneNumber(s_phone + irow * byte_width, tld.rng, s_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ACCTBAL(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_ACCTBAL].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_ACCTBAL));
+ Decimal128* s_acctbal = reinterpret_cast<Decimal128*>(
+ tld.batch[SUPPLIER::S_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_acctbal[irow] = {dist(tld.rng)};
+ }
+ return Status::OK();
+ }
+
+ Status S_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[SUPPLIER::S_COMMENT],
+ g_text.GenerateComments(tld.to_generate, 25, 100, tld.rng));
+ ModifyComments(thread_index, "Recommends", good_rows_);
+ ModifyComments(thread_index, "Complaints", bad_rows_);
+ }
+ return Status::OK();
+ }
+
+ void ModifyComments(size_t thread_index, const char* review,
+ const std::vector<int64_t>& indices) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ const int32_t* offsets = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[1]->data());
+ char* str = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[2]->mutable_data());
+ const char* customer = "Customer";
+ const int32_t customer_length = static_cast<int32_t>(std::strlen(customer));
+ const int32_t review_length = static_cast<int32_t>(std::strlen(review));
+
+ auto it = std::lower_bound(indices.begin(), indices.end(), tld.suppkey_start);
+ for (; it != indices.end() && *it < tld.suppkey_start + tld.to_generate; it++) {
+ int64_t idx_in_batch = *it - tld.suppkey_start;
+ char* out = str + offsets[idx_in_batch];
+ int32_t str_length = offsets[idx_in_batch + 1] - offsets[idx_in_batch];
+ std::uniform_int_distribution<int32_t> gap_dist(
+ 0, str_length - customer_length - review_length);
+ int32_t gap = gap_dist(tld.rng);
+ int32_t total_length = customer_length + gap + review_length;
+ std::uniform_int_distribution<int32_t> start_dist(0, str_length - total_length);
+ int32_t start = start_dist(tld.rng);
+ std::memcpy(out + start, customer, customer_length);
+ std::memcpy(out + start + customer_length + gap, review, review_length);
+ }
+ }
+
+ struct ThreadLocalData {
+ random::pcg32_fast rng;
+ int64_t suppkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+ std::vector<int64_t> good_rows_;
+ std::vector<int64_t> bad_rows_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_;
+ std::atomic<int64_t> rows_generated_;
+ double scale_factor_;
+ int64_t batch_size_;
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class PartGenerator : public TpchTableGenerator {
+ public:
+ explicit PartGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, double scale_factor, int64_t batch_size,
+ int64_t seed) override {
+ seed_ = seed;
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetPartOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_, seed_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch, gen_->NextPartBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->part_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ double scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class PartSuppGenerator : public TpchTableGenerator {
+ public:
+ explicit PartSuppGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, double scale_factor, int64_t batch_size,
+ int64_t seed) override {
+ seed_ = seed;
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetPartSuppOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_, seed_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextPartSuppBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->partsupp_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ double scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class CustomerGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, double scale_factor, int64_t batch_size,
+ int64_t seed) override {
+ seed_ = seed;
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, gen_list_));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ thread_local_data_.resize(num_threads);
+ random::pcg64_fast seed_rng(seed_);
+ for (ThreadLocalData& tld : thread_local_data_) tld.rng.seed(kSeedDist(seed_rng));
+
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(C_CUSTKEY) \
+ F(C_NAME) \
+ F(C_ADDRESS) \
+ F(C_NATIONKEY) \
+ F(C_PHONE) \
+ F(C_ACCTBAL) \
+ F(C_MKTSEGMENT) \
+ F(C_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct CUSTOMER {
+ enum {
+ FOR_EACH_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) {#col, CUSTOMER::col},
+ const std::unordered_map<std::string, int> kNameMap = {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kGenerators = {FOR_EACH_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(),
+ utf8(),
+ utf8(),
+ int32(),
+ fixed_size_binary(15),
+ decimal(12, 2),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index) {
+ if (done_.load()) return Status::OK();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.custkey_start = rows_generated_.fetch_add(batch_size_);
+ if (tld.custkey_start >= rows_to_generate_) return Status::OK();
+
+ tld.to_generate = std::min(batch_size_, rows_to_generate_ - tld.custkey_start);
+
+ tld.batch.resize(CUSTOMER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for (int col : gen_list_) RETURN_NOT_OK(kGenerators[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for (size_t i = 0; i < gen_list_.size(); i++) {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_generated_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_generated_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if (is_last_batch) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ return Status::OK();
+ }
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(kTypes[column], tld.to_generate, {nullptr, std::move(buff)});
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status C_CUSTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_CUSTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_CUSTKEY));
+ int32_t* c_custkey = reinterpret_cast<int32_t*>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ c_custkey[irow] = static_cast<int32_t>(tld.custkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_NAME].kind() == Datum::NONE) {
+ RETURN_NOT_OK(C_CUSTKEY(thread_index));
+ const int32_t* c_custkey = reinterpret_cast<const int32_t*>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->data());
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ const char* customer = "Customer#";
+ const size_t customer_length = std::strlen(customer);
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ int num_digits = GetNumDigits(c_custkey[irow]);
+ int num_chars = std::max(num_digits, 9);
+ offsets[irow + 1] =
+ static_cast<int32_t>(offsets[irow] + num_chars + customer_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[tld.to_generate]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = str + offsets[irow];
+ std::memcpy(out, customer, customer_length);
+ AppendNumberPaddedToNineDigits(out + customer_length, c_custkey[irow]);
+ }
+ ArrayData ad(utf8(), tld.to_generate,
+ {nullptr, std::move(offset_buff), std::move(str_buff)});
+ tld.batch[CUSTOMER::C_NAME] = std::move(ad);
+ }
+ return Status::OK();
+ }
+
+ Status C_ADDRESS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_ADDRESS].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[CUSTOMER::C_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status C_NATIONKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_NATIONKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t* c_nationkey = reinterpret_cast<int32_t*>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status C_PHONE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_PHONE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(C_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[CUSTOMER::C_PHONE]);
+ const int32_t* c_nationkey = reinterpret_cast<const int32_t*>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->data());
+ char* c_phone = reinterpret_cast<char*>(
+ tld.batch[CUSTOMER::C_PHONE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ GeneratePhoneNumber(c_phone + irow * byte_width, tld.rng, c_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_ACCTBAL(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_ACCTBAL].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_ACCTBAL));
+ Decimal128* c_acctbal = reinterpret_cast<Decimal128*>(
+ tld.batch[CUSTOMER::C_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_acctbal[irow] = {dist(tld.rng)};
+ }
+ return Status::OK();
+ }
+
+ Status C_MKTSEGMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_MKTSEGMENT].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_MKTSEGMENT));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[CUSTOMER::C_MKTSEGMENT]);
+ char* c_mktsegment = reinterpret_cast<char*>(
+ tld.batch[CUSTOMER::C_MKTSEGMENT].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int32_t> dist(0, kNumSegments - 1);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = c_mktsegment + irow * byte_width;
+ int str_idx = dist(tld.rng);
+ std::strncpy(out, Segments[str_idx], byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[CUSTOMER::C_COMMENT],
+ g_text.GenerateComments(tld.to_generate, 29, 116, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ random::pcg32_fast rng;
+ int64_t custkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_{0};
+ std::atomic<int64_t> rows_generated_ = {0};
+ double scale_factor_{0};
+ int64_t batch_size_{0};
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class OrdersGenerator : public TpchTableGenerator {
+ public:
+ explicit OrdersGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, double scale_factor, int64_t batch_size,
+ int64_t seed) override {
+ seed_ = seed;
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetOrdersOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_, seed_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch, gen_->NextOrdersBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->orders_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ double scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class LineitemGenerator : public TpchTableGenerator {
+ public:
+ explicit LineitemGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, double scale_factor, int64_t batch_size,
+ int64_t seed) override {
+ seed_ = seed;
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetLineItemOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_, seed_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextLineItemBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->lineitem_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ double scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class NationGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, double /*scale_factor*/,
+ int64_t /*batch_size*/, int64_t seed) override {
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, column_indices_));
+ seed_ = seed;
+ rng_.seed(seed_);
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t /*num_threads*/, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback /*schedule_task_callback*/) override {
+ std::shared_ptr<Buffer> N_NATIONKEY_buffer =
+ Buffer::Wrap(kNationKey, sizeof(kNationKey));
+ ArrayData N_NATIONKEY_arraydata(int32(), kRowCount,
+ {nullptr, std::move(N_NATIONKEY_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> N_NAME_buffer,
+ AllocateBuffer(kRowCount * kNameByteWidth));
+ char* N_NAME = reinterpret_cast<char*>(N_NAME_buffer->mutable_data());
+ for (size_t i = 0; i < kRowCount; i++)
+ std::strncpy(N_NAME + kNameByteWidth * i, kCountryNames[i], kNameByteWidth);
+ ArrayData N_NAME_arraydata(fixed_size_binary(kNameByteWidth), kRowCount,
+ {nullptr, std::move(N_NAME_buffer)});
+
+ std::shared_ptr<Buffer> N_REGIONKEY_buffer =
+ Buffer::Wrap(kRegionKey, sizeof(kRegionKey));
+ ArrayData N_REGIONKEY_arraydata(int32(), kRowCount,
+ {nullptr, std::move(N_REGIONKEY_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(Datum N_COMMENT_datum,
+ g_text.GenerateComments(kRowCount, 31, 114, rng_));
+
+ std::vector<Datum> fields = {
+ std::move(N_NATIONKEY_arraydata), std::move(N_NAME_arraydata),
+ std::move(N_REGIONKEY_arraydata), std::move(N_COMMENT_datum)};
+
+ std::vector<Datum> result;
+ for (const int& col : column_indices_) result.push_back(fields[col]);
+ ARROW_ASSIGN_OR_RAISE(ExecBatch batch, ExecBatch::Make(std::move(result)));
+ output_callback(std::move(batch));
+ finished_callback(static_cast<int64_t>(1));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ random::pcg32_fast rng_;
+
+ static constexpr size_t kRowCount = 25;
+ static constexpr int32_t kNameByteWidth = 25;
+ static constexpr int32_t kNationKey[kRowCount] = {0, 1, 2, 3, 4, 5, 6, 7, 8,
+ 9, 10, 11, 12, 13, 14, 15, 16, 17,
+ 18, 19, 20, 21, 22, 23, 24};
+ static constexpr const char* kCountryNames[kRowCount] = {
+ "ALGERIA", "ARGENTINA", "BRAZIL", "CANADA", "EGYPT",
+ "ETHIOPIA", "FRANCE", "GERMANY", "INDIA", "INDONESIA",
+ "IRAN", "IRAQ", "JAPAN", "JORDAN", "KENYA",
+ "MOROCCO", "MOZAMBIQUE", "PERU", "CHINA", "ROMANIA",
+ "SAUDI ARABIA", "VIETNAM", "RUSSIA", "UNITED KINGDOM", "UNITED STATES"};
+ static constexpr int32_t kRegionKey[kRowCount] = {0, 1, 1, 1, 4, 0, 3, 3, 2, 2, 4, 4, 2,
+ 4, 0, 0, 0, 1, 2, 3, 4, 2, 3, 3, 1};
+
+ struct NATION {
+ enum {
+ N_NATIONKEY,
+ N_NAME,
+ N_REGIONKEY,
+ N_COMMENT,
+ };
+ };
+
+ const std::unordered_map<std::string, int> kNameMap = {
+ {"N_NATIONKEY", NATION::N_NATIONKEY},
+ {"N_NAME", NATION::N_NAME},
+ {"N_REGIONKEY", NATION::N_REGIONKEY},
+ {"N_COMMENT", NATION::N_COMMENT},
+ };
+
+ std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(),
+ fixed_size_binary(kNameByteWidth),
+ int32(),
+ utf8(),
+ };
+
+ std::shared_ptr<Schema> schema_;
+ std::vector<int> column_indices_;
+};
+
+constexpr int32_t NationGenerator::kNationKey[NationGenerator::kRowCount];
+constexpr const char* NationGenerator::kCountryNames[NationGenerator::kRowCount];
+constexpr int32_t NationGenerator::kRegionKey[NationGenerator::kRowCount];
+
+class RegionGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, double /*scale_factor*/,
+ int64_t /*batch_size*/, int64_t seed) override {
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, column_indices_));
+ seed_ = seed;
+ rng_.seed(seed_);
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback /*schedule_task_callback*/) override {
+ std::shared_ptr<Buffer> R_REGIONKEY_buffer =
+ Buffer::Wrap(kRegionKey, sizeof(kRegionKey));
+ ArrayData R_REGIONKEY_arraydata(int32(), kRowCount,
+ {nullptr, std::move(R_REGIONKEY_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> R_NAME_buffer,
+ AllocateBuffer(kRowCount * kNameByteWidth));
+ char* R_NAME_data = reinterpret_cast<char*>(R_NAME_buffer->mutable_data());
+ for (size_t i = 0; i < kRowCount; i++)
+ std::strncpy(R_NAME_data + kNameByteWidth * i, kRegionNames[i], kNameByteWidth);
+ ArrayData R_NAME_arraydata(kTypes[static_cast<int>(REGION::R_NAME)], kRowCount,
+ {nullptr, std::move(R_NAME_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(Datum R_COMMENT_datum,
+ g_text.GenerateComments(kRowCount, 31, 115, rng_));
+
+ std::vector<Datum> fields = {std::move(R_REGIONKEY_arraydata),
+ std::move(R_NAME_arraydata), std::move(R_COMMENT_datum)};
+ std::vector<Datum> result;
+ for (const int& col : column_indices_) result.push_back(fields[col]);
+ ARROW_ASSIGN_OR_RAISE(ExecBatch batch, ExecBatch::Make(std::move(result)));
+ output_callback(std::move(batch));
+ finished_callback(static_cast<int64_t>(1));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ random::pcg32_fast rng_;
+
+ static constexpr size_t kRowCount = 5;
+ static constexpr int32_t kNameByteWidth = 25;
+ static constexpr int32_t kRegionKey[kRowCount] = {0, 1, 2, 3, 4};
+ static constexpr const char* kRegionNames[kRowCount] = {"AFRICA", "AMERICA", "ASIA",
+ "EUROPE", "MIDDLE EAST"};
+
+ struct REGION {
+ enum {
+ R_REGIONKEY,
+ R_NAME,
+ R_COMMENT,
+ kNumColumns,
+ };
+ };
+
+ const std::unordered_map<std::string, int> kNameMap = {
+ {"R_REGIONKEY", REGION::R_REGIONKEY},
+ {"R_NAME", REGION::R_NAME},
+ {"R_COMMENT", REGION::R_COMMENT},
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(),
+ fixed_size_binary(kNameByteWidth),
+ utf8(),
+ };
+
+ std::shared_ptr<Schema> schema_;
+ std::vector<int> column_indices_;
+};
+
+constexpr int32_t RegionGenerator::kRegionKey[RegionGenerator::kRowCount];
+constexpr const char* RegionGenerator::kRegionNames[RegionGenerator::kRowCount];
+
+class TpchNode : public ExecNode {
+ public:
+ TpchNode(ExecPlan* plan, const char* name,
+ std::unique_ptr<TpchTableGenerator> generator)
+ : ExecNode(plan, {}, {}, generator->schema(), /*num_outputs=*/1),
+ name_(name),
+ generator_(std::move(generator)) {}
+
+ const char* kind_name() const override { return name_; }
+
+ [[noreturn]] static void NoInputs() {
+ Unreachable("TPC-H node should never have any inputs");
+ }
+
+ [[noreturn]] void InputReceived(ExecNode*, ExecBatch) override { NoInputs(); }
+
+ [[noreturn]] void ErrorReceived(ExecNode*, Status) override { NoInputs(); }
+
+ [[noreturn]] void InputFinished(ExecNode*, int) override { NoInputs(); }
+
+ Status StartProducing() override {
+ return generator_->StartProducing(
+ thread_indexer_.Capacity(),
+ [this](ExecBatch batch) { this->OutputBatchCallback(std::move(batch)); },
+ [this](int64_t num_batches) { this->FinishedCallback(num_batches); },
+ [this](std::function<Status(size_t)> func) -> Status {
+ return this->ScheduleTaskCallback(std::move(func));
+ });
+ }
+
+ void PauseProducing(ExecNode* output) override {}
+ void ResumeProducing(ExecNode* output) override {}
+
+ void StopProducing(ExecNode* output) override {
+ DCHECK_EQ(output, outputs_[0]);
+ StopProducing();
+ }
+
+ void StopProducing() override {
+ if (generator_->Abort()) std::ignore = task_group_.End();
+ }
+
+ Future<> finished() override { return task_group_.OnFinished(); }
+
+ private:
+ void OutputBatchCallback(ExecBatch batch) {
+ outputs_[0]->InputReceived(this, std::move(batch));
+ }
+
+ void FinishedCallback(int64_t total_num_batches) {
+ outputs_[0]->InputFinished(this, static_cast<int>(total_num_batches));
+ std::ignore = task_group_.End();
+ }
+
+ Status ScheduleTaskCallback(std::function<Status(size_t)> func) {
+ auto executor = plan_->exec_context()->executor();
+ if (executor) {
+ RETURN_NOT_OK(task_group_.AddTaskIfNotEnded([&] {
Review comment:
Added.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r837766413
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -42,40 +42,23 @@ class ARROW_EXPORT TpchGen {
* table from that single TpchGen instance. Note: Every batch will be scheduled as a new
* task using the ExecPlan's scheduler.
*/
- static Result<TpchGen> Make(ExecPlan* plan, double scale_factor = 1.0,
- int64_t batch_size = 4096,
- util::optional<int64_t> seed = util::nullopt);
+ static Result<std::unique_ptr<TpchGen>> Make(
Review comment:
As the commit title suggested, the point here is to hide implementation details from the `.h`, and minimize header inclusion.
It can be a `unique_ptr`, a `shared_ptr` or the pimpl pattern, but any of these patterns require some form of dynamic allocation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082144362
Hmm, there's ASAN issue that needs to be looked at before merging this. I can do that tomorrow if you want.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1057461689
> @jonkeane They get generated at the same time if you use the same TpchGen object.
Cool, so then the next thing would be to figure out if there's a way to emit and write these record batch readers simultaneously. (and that might actually need to to be done all in C++, the R execution paradigm might make it hard to have one plan that is emitting two sets of record batches and writes them simultaneously)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r818746754
##########
File path: r/src/compute-exec.cpp
##########
@@ -277,4 +287,135 @@ std::shared_ptr<compute::ExecNode> ExecNode_ReadFromRecordBatchReader(
return MakeExecNodeOrStop("source", plan.get(), {}, options);
}
+// [[arrow::export]]
+std::shared_ptr<arrow::RecordBatchReader> Tpch_Dbgen(
+ const std::shared_ptr<compute::ExecPlan>& plan,
+ int scale_factor,
+ std::string table_name
+ ) {
+
+ auto gen = ValueOrStop(arrow::compute::TpchGen::Make(plan.get(), scale_factor));
+
+ compute::ExecNode *table;
+ if (table_name == "part") {
+ table = ValueOrStop(gen.Part());
+ } else if (table_name == "supplier") {
+ table = ValueOrStop(gen.Supplier());
+ } else if (table_name == "partsupp") {
+ table = ValueOrStop(gen.PartSupp());
+ } else if (table_name == "customer") {
+ table = ValueOrStop(gen.Customer());
+ } else if (table_name == "nation") {
+ table = ValueOrStop(gen.Nation());
+ } else if (table_name == "lineitem") {
+ table = ValueOrStop(gen.Lineitem());
+ } else if (table_name == "region") {
+ table = ValueOrStop(gen.Region());
+ } else if (table_name == "orders") {
+ table = ValueOrStop(gen.Orders());
+ } else {
+ cpp11::stop("That's not a valid table name");
+ }
+
+ arrow::AsyncGenerator<arrow::util::optional<compute::ExecBatch>> sink_gen;
+
+ MakeExecNodeOrStop("sink", plan.get(), {table},
+ compute::SinkNodeOptions{&sink_gen});
+
+ StopIfNotOk(plan->Validate());
+ StopIfNotOk(plan->StartProducing());
+
+ // If the generator is destroyed before being completely drained, inform plan
+ std::shared_ptr<void> stop_producing{nullptr, [plan](...) {
+ bool not_finished_yet =
+ plan->finished().TryAddCallback([&plan] {
+ return [plan](const arrow::Status&) {};
+ });
+
+ if (not_finished_yet) {
+ plan->StopProducing();
+ }
+ }};
+
+ return compute::MakeGeneratorReader(
+ table->output_schema(),
+ [stop_producing, plan, sink_gen] { return sink_gen(); }, gc_memory_pool());
+}
+
+// [[arrow::export]]
+void Tpch_Dbgen_Write(
+ const std::shared_ptr<compute::ExecPlan>& plan,
+ int scale_factor,
+ std::string table_name,
+ const std::shared_ptr<fs::FileSystem>& filesystem, std::string base_dir,
+ arrow::dataset::ExistingDataBehavior existing_data_behavior, int max_partitions
+) {
+ auto gen = ValueOrStop(arrow::compute::TpchGen::Make(plan.get(), scale_factor));
+
+ compute::ExecNode *table;
+ if (table_name == "part") {
+ table = ValueOrStop(gen.Part());
+ } else if (table_name == "supplier") {
+ table = ValueOrStop(gen.Supplier());
+ } else if (table_name == "partsupp") {
+ table = ValueOrStop(gen.PartSupp());
+ } else if (table_name == "customer") {
+ table = ValueOrStop(gen.Customer());
+ } else if (table_name == "nation") {
+ table = ValueOrStop(gen.Nation());
+ } else if (table_name == "lineitem") {
+ table = ValueOrStop(gen.Lineitem());
+ } else if (table_name == "region") {
+ table = ValueOrStop(gen.Region());
+ } else if (table_name == "orders") {
+ table = ValueOrStop(gen.Orders());
+ } else {
+ cpp11::stop("That's not a valid table name");
+ }
+
+ // TODO: unhardcode this once it's working
+ auto base_path = base_dir + "/parquet_dataset";
+ filesystem->CreateDir(base_path);
+
+ auto format = std::make_shared<ds::ParquetFileFormat>();
+
+ ds::FileSystemDatasetWriteOptions write_options;
+ write_options.file_write_options = format->DefaultWriteOptions();
+ write_options.existing_data_behavior = ds::ExistingDataBehavior::kDeleteMatchingPartitions;
+ write_options.filesystem = filesystem;
+ write_options.base_dir = base_path;
+ write_options.partitioning = arrow::dataset::Partitioning::Default();
+ write_options.basename_template = "part{i}.parquet";
+ write_options.max_partitions = 1024;
Review comment:
A lot of this is hard coded for now to get it working — one thing that surprise me a little bit is that we have slightly different write options for this than we do for datasets. I think I need to wire up `FileSystemDatasetWriteOptions` (or figure out if it gets translated in the right way to get to that point with our R wiring of `FileWriteOptions`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r825006578
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
Review comment:
Yeah, `gen_list` is a list of indices of columns. So if we e.g. wanted to project the same column twice or change the order, we rearrange the gen list.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r825197134
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string> &columns,
+ const std::vector<std::shared_ptr<DataType>> &types,
+ const std::unordered_map<std::string, int> &name_map,
+ std::vector<int> &gen_list)
+ {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if(columns.empty())
+ {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for(auto pair : name_map)
+ {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ }
+ else
+ {
+ for(const std::string &col : columns)
+ {
+ auto entry = name_map.find(col);
+ if(entry == name_map.end())
+ return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+ }
+
+ static TpchText g_text;
+
+ Status TpchText::InitIfNeeded(random::pcg32_fast &rng)
+ {
+ if(done_.load())
+ return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(!text_)
+ {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char *out = reinterpret_cast<char *>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+ while(done_.load() == false)
+ {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while(GenerateSentence(try_offset, rng, temp_buff))
+ known_valid_offset = try_offset;
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(done_.load())
+ return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ int64_t memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ std::memcpy(out + generated_offset_, temp_buff, memcpy_size);
+ generated_offset_ += memcpy_size;
+ if(generated_offset_ == kTextBytes)
+ done_.store(true);
+ }
+ }
+ return Status::OK();
+ }
+
+ Result<Datum> TpchText::GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng)
+ {
+ RETURN_NOT_OK(InitIfNeeded(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer, AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for(size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + length_dist(rng);
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer, AllocateBuffer(offsets[num_comments]));
+ char *comments = reinterpret_cast<char *>(comment_buffer->mutable_data());
+ for(size_t i = 0; i < num_comments; i++)
+ {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments, { nullptr, std::move(offset_buffer), std::move(comment_buffer) });
+ return std::move(ad);
+ }
+
+ Result<Datum> RandomVString(
+ random::pcg32_fast &rng,
+ int64_t num_rows,
+ int32_t min_length,
+ int32_t max_length)
+ {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff, AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for(int64_t i = 1; i <= num_rows; i++)
+ offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff, AllocateBuffer(offsets[num_rows]));
+ char *str = reinterpret_cast<char *>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for(int32_t i = 0; i < offsets[num_rows]; i++)
+ str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, { nullptr, std::move(offset_buff), std::move(str_buff) });
+ return std::move(ad);
+ }
+
+ void AppendNumber(char *&out, int num_digits, int32_t x)
+ {
+ out += (num_digits - 1);
+ while(x > 0)
+ {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+ }
+
+ void GeneratePhoneNumber(
+ char *out,
+ random::pcg32_fast &rng,
+ int32_t country)
+ {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+ }
+
+ static constexpr uint32_t STARTDATE = 8035; // January 1, 1992 is 8035 days after January 1, 1970
+ static constexpr uint32_t CURRENTDATE = 9298; // June 17, 1995 is 9298 days after January 1, 1970
+ static constexpr uint32_t ENDDATE = 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+ const char *NameParts[] =
+ {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque", "black", "blanched", "blue",
+ "blush", "brown", "burlywood", "burnished", "chartreuse", "chiffon", "chocolate", "coral",
+ "cornflower", "cornsilk", "cream", "cyan", "dark", "deep", "dim", "dodger", "drab", "firebrick",
+ "floral", "forest", "frosted", "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender", "lawn", "lemon", "light", "lime", "linen",
+ "magenta", "maroon", "medium", "metallic", "midnight", "mint", "misty", "moccasin", "navajo",
+ "navy", "olive", "orange", "orchid", "pale", "papaya", "peach", "peru", "pink", "plum", "powder",
+ "puff", "purple", "red", "rose", "rosy", "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel", "tan", "thistle", "tomato", "turquoise", "violet",
+ "wheat", "white", "yellow",
+ };
+ static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+ const char *Types_1[] =
+ {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+ };
+ static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+ const char *Types_2[] =
+ {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+ };
+ static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+ const char *Types_3[] =
+ {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+ };
+ static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+ const char *Containers_1[] =
+ {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+ };
+ static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+ const char *Containers_2[] =
+ {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+ };
+ static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+ const char *Segments[] =
+ {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+ };
+ static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+ const char *Priorities[] =
+ {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+ };
+ static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+ const char *Instructions[] =
+ {
+ "DELIVER IN PERSON", "COLLECT COD", "NONE", "TAKE BACK RETURN",
+ };
+ static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+ const char *Modes[] =
+ {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+ };
+ static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+ const char *Nouns[] =
+ {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ", "dependencies ", "excuses ",
+ "platelets ", "asymptotes ", "courts ", "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ", "forges ", "braids ", "hockey players ", "frays ",
+ "warhorses ", "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ", "waters ", "orbits ", "gifts ", "sheaves ",
+ "depths ", "sentiments ", "decoys ", "realms ", "pains ", "grouches ", "escapades ",
+ };
+ static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+ const char *Verbs[] =
+ {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ", "boost ", "affix ", "detect ", "integrate ",
+ "maintain ", "nod ", "was ", "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ", "poach ", "serve ", "run ", "dazzle ",
+ "snooze ", "doze ", "unwind ", "kindle ", "play ", "hang ", "believe ", "doubt ",
+ };
+ static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+ const char *Adjectives[] =
+ {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ", "quiet ", "ruthless ", "thin ",
+ "close ", "dogged ", "daring ", "brave ", "stealthy ", "permanent ", "enticing ", "idle ", "busy ",
+ "regular ", "final ", "ironic ", "even ", "bold ", "silent ",
+ };
+ static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+ const char *Adverbs[] =
+ {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ", "blithely ", "quickly ", "fluffily ",
+ "slowly ", "quietly ", "ruthlessly ", "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ", "ironically ", "evenly ", "boldly ",
+ "silently ",
+ };
+ static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+ const char *Prepositions[] =
+ {
+ "about ", "above ", "according to ", "across ", "after ", "against ", "along ", "alongside of ", "among ",
+ "around ", "at ", "atop ", "before ", "behind ", "beneath ", "beside ", "besides ", "between ", "beyond ",
+ "beyond ", "by ", "despite ", "during ", "except ", "for ", "from ", "in place of ", "inside ", "instead of ",
+ "into ", "near ", "of ", "on ", "outside ", "over ", "past ", "since ", "through ", "throughout ", "to ",
+ "toward ", "under ", "until ", "up ", "upon ", "without ", "with ", "within ",
+ };
+ static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+ const char *Auxiliaries[] =
+ {
+ "do ", "may ", "might ", "shall ", "will ", "would ", "can ", "could ", "should ", "ought to ", "must ",
+ "will have to ", "shall have to ", "could have to ", "should have to ", "must have to ", "need to ", "try to ",
+ };
+ static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+ const char *Terminators[] =
+ {
+ ".", ";", ":", "?", "!", "--",
+ };
+ static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+ bool TpchText::GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices)
+ {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char *word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if(offset + length > kChunkSize)
+ return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+ }
+
+ bool TpchText::GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+ }
+
+ bool TpchText::GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+ }
+
+ bool TpchText::GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+ }
+
+ bool TpchText::GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+ }
+
+ bool TpchText::GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+ }
+
+ bool TpchText::GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+ }
+
+ bool TpchText::GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if(result)
+ std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+ }
+
+ bool TpchText::GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char *comma_space = ", ";
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ GenerateAdverb(offset, rng, arr);
+ GenerateAdjective(offset, rng, arr);
+ GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ bool TpchText::GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ bool TpchText::GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ const char *the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+ }
+
+ bool TpchText::GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 5:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ using GenerateColumnFn = std::function<Status(size_t)>;
+ class PartAndPartSupplierGenerator
+ {
+ public:
+ Status Init(
+ size_t num_threads,
+ int64_t batch_size,
+ float scale_factor)
+ {
+ if(!inited_)
+ {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ thread_local_data_.resize(num_threads);
+ for(ThreadLocalData &tld : thread_local_data_)
+ {
+ // 5 is the maximum number of different strings we need to concatenate
+ tld.string_indices.resize(5 * batch_size_);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const
+ {
+ return part_batches_generated_.load();
+ }
+
+ int64_t partsupp_batches_generated() const
+ {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, part_types_, part_name_map_, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, partsupp_types_, partsupp_name_map_, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch()
+ {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ bool all_generated = part_rows_generated_ == part_rows_to_generate_;
+ if(!part_output_queue_.empty())
+ {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ }
+ else if(all_generated)
+ {
+ return util::nullopt;
+ }
+ else
+ {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate = std::min(
+ batch_size_,
+ part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for(int col : part_cols_)
+ {
+ RETURN_NOT_OK(part_generators_[col](thread_index));
+ }
+ for(int col : partsupp_cols_)
+ RETURN_NOT_OK(partsupp_generators_[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for(size_t i = 0; i < part_cols_.size(); i++)
+ {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if(!partsupp_cols_.empty())
+ {
+ std::vector<ExecBatch> partsupp_results;
+ for(size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++)
+ {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for(size_t icol = 0; icol < partsupp_cols_.size(); icol++)
+ {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for(ExecBatch &eb : partsupp_results)
+ {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch()
+ {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ if(!partsupp_output_queue_.empty())
+ {
+ ExecBatch result = std::move(partsupp_output_queue_.front());
+ partsupp_output_queue_.pop();
+ return std::move(result);
+ }
+ }
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if(part_rows_generated_ == part_rows_to_generate_)
+ {
+ return util::nullopt;
+ }
+ else
+ {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate = std::min(
+ batch_size_,
+ part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for(int col : part_cols_)
+ RETURN_NOT_OK(part_generators_[col](thread_index));
+ for(int col : partsupp_cols_)
+ RETURN_NOT_OK(partsupp_generators_[col](thread_index));
+ if(!part_cols_.empty())
+ {
+ std::vector<Datum> part_result(part_cols_.size());
+ for(size_t i = 0; i < part_cols_.size(); i++)
+ {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch part_batch, ExecBatch::Make(std::move(part_result)));
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ part_output_queue_.emplace(std::move(part_batch));
+ }
+ }
+ std::vector<ExecBatch> partsupp_results;
+ for(size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++)
+ {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for(size_t icol = 0; icol < partsupp_cols_.size(); icol++)
+ {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ for(size_t i = 1; i < partsupp_results.size(); i++)
+ partsupp_output_queue_.emplace(std::move(partsupp_results[i]));
+ }
+ return std::move(partsupp_results[0]);
+ }
+
+ private:
+#define FOR_EACH_PART_COLUMN(F) \
+ F(P_PARTKEY) \
+ F(P_NAME) \
+ F(P_MFGR) \
+ F(P_BRAND) \
+ F(P_TYPE) \
+ F(P_SIZE) \
+ F(P_CONTAINER) \
+ F(P_RETAILPRICE) \
+ F(P_COMMENT)
+
+#define FOR_EACH_PARTSUPP_COLUMN(F) \
+ F(PS_PARTKEY) \
+ F(PS_SUPPKEY) \
+ F(PS_AVAILQTY) \
+ F(PS_SUPPLYCOST) \
+ F(PS_COMMENT) \
+
+#define MAKE_ENUM(col) col,
+ struct PART
+ {
+ enum
+ {
+ FOR_EACH_PART_COLUMN(MAKE_ENUM)
+ kNumCols,
+ };
+ };
+ struct PARTSUPP
+ {
+ enum
+ {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_ENUM)
+ kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) \
+ { #col, PART::col },
+ const std::unordered_map<std::string, int> part_name_map_ =
+ {
+ FOR_EACH_PART_COLUMN(MAKE_STRING_MAP)
+ };
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) \
+ { #col, PARTSUPP::col },
+ const std::unordered_map<std::string, int> partsupp_name_map_ =
+ {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_STRING_MAP)
+ };
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> part_generators_ =
+ {
+ FOR_EACH_PART_COLUMN(MAKE_FN_ARRAY)
+ };
+ std::vector<GenerateColumnFn> partsupp_generators_ =
+ {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_FN_ARRAY)
+ };
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> part_types_ =
+ {
+ int32(),
+ utf8(),
+ fixed_size_binary(25),
+ fixed_size_binary(10),
+ utf8(),
+ int32(),
+ fixed_size_binary(10),
+ decimal(12, 2),
+ utf8(),
+ };
+
+ const std::vector<std::shared_ptr<DataType>> partsupp_types_ =
+ {
+ int32(),
+ int32(),
+ int32(),
+ decimal(12, 2),
+ utf8(),
+ };
+
+ Status AllocatePartBatch(size_t thread_index, int column)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.part[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*part_types_[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff, AllocateBuffer(tld.part_to_generate * byte_width));
+ ArrayData ad(part_types_[column], tld.part_to_generate, { nullptr, std::move(buff) });
+ tld.part[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status P_PARTKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_PARTKEY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_PARTKEY));
+ int32_t *p_partkey = reinterpret_cast<int32_t *>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.part_to_generate; i++)
+ {
+ p_partkey[i] = (tld.partkey_start + i + 1);
+ ARROW_DCHECK(1 <= p_partkey[i] && p_partkey[i] <= part_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_NAME(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_NAME].kind() == Datum::NONE)
+ {
+ std::uniform_int_distribution<uint8_t> dist(0, static_cast<uint8_t>(kNumNameParts - 1));
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff, AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ size_t string_length = 0;
+ for(int ipart = 0; ipart < 5; ipart++)
+ {
+ uint8_t name_part_index = dist(tld.rng);
+ tld.string_indices[irow * 5 + ipart] = name_part_index;
+ string_length += std::strlen(NameParts[name_part_index]);
+ }
+ // Add 4 because there is a space between each word (i.e. four spaces)
+ offsets[irow + 1] = offsets[irow] + string_length + 4;
+ }
+ // Add an extra byte for the space after in the very last string.
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer, AllocateBuffer(offsets[tld.part_to_generate] + 1));
+ char *strings = reinterpret_cast<char *>(string_buffer->mutable_data());
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ char *row = strings + offsets[irow];
+ for(int ipart = 0; ipart < 5; ipart++)
+ {
+ uint8_t name_part_index = tld.string_indices[irow * 5 + ipart];
+ const char *part = NameParts[name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ *row++ = ' ';
+ }
+ }
+ ArrayData ad(part_types_[PART::P_NAME], tld.part_to_generate, { nullptr, std::move(offset_buff), std::move(string_buffer) });
+ Datum datum(ad);
+ tld.part[PART::P_NAME] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_MFGR(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_MFGR].kind() == Datum::NONE)
+ {
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char *manufacturer = "Manufacturer#";
+ const size_t manufacturer_length = std::strlen(manufacturer);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_MFGR));
+ char *p_mfgr = reinterpret_cast<char *>(tld.part[PART::P_MFGR].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*part_types_[PART::P_MFGR]);
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ std::strncpy(p_mfgr + byte_width * irow, manufacturer, byte_width);
+ char mfgr_id = '0' + dist(tld.rng);
+ *(p_mfgr + byte_width * irow + manufacturer_length) = mfgr_id;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_BRAND(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_BRAND].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(P_MFGR(thread_index));
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char *brand = "Brand#";
+ const size_t brand_length = std::strlen(brand);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_BRAND));
+ const char *p_mfgr = reinterpret_cast<const char *>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->data());
+ char *p_brand = reinterpret_cast<char *>(
+ tld.part[PART::P_BRAND].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*part_types_[PART::P_BRAND]);
+ int32_t mfgr_byte_width = arrow::internal::GetByteWidth(*part_types_[PART::P_MFGR]);
+ const size_t mfgr_id_offset = std::strlen("Manufacturer#");
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ char *row = p_brand + byte_width * irow;
+ char mfgr_id = *(p_mfgr + irow * mfgr_byte_width + mfgr_id_offset);
+ char brand_id = '0' + dist(tld.rng);
+ std::strncpy(row, brand, byte_width);
+ *(row + brand_length) = mfgr_id;
+ *(row + brand_length + 1) = brand_id;
+ irow += 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_TYPE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_TYPE].kind() == Datum::NONE)
+ {
+ using D = std::uniform_int_distribution<uint8_t>;
+ D dists[] =
+ {
+ D{ 0, static_cast<uint8_t>(kNumTypes_1 - 1) },
+ D{ 0, static_cast<uint8_t>(kNumTypes_2 - 1) },
+ D{ 0, static_cast<uint8_t>(kNumTypes_3 - 1) },
+ };
+
+ const char **types[] = { Types_1, Types_2, Types_3 };
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff, AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ size_t string_length = 0;
+ for(int ipart = 0; ipart < 3; ipart++)
+ {
+ uint8_t name_part_index = dists[ipart](tld.rng);
+ tld.string_indices[irow * 3 + ipart] = name_part_index;
+ string_length += std::strlen(types[ipart][name_part_index]);
+ }
+ offsets[irow + 1] = offsets[irow] + string_length;
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer, AllocateBuffer(offsets[tld.part_to_generate]));
+ char *strings = reinterpret_cast<char *>(string_buffer->mutable_data());
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ char *row = strings + offsets[irow];
+ for(int ipart = 0; ipart < 3; ipart++)
+ {
+ uint8_t name_part_index = tld.string_indices[irow * 3 + ipart];
+ const char *part = types[ipart][name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ }
+ }
+ ArrayData ad(part_types_[PART::P_TYPE], tld.part_to_generate, { nullptr, std::move(offset_buff), std::move(string_buffer) });
+ Datum datum(ad);
+ tld.part[PART::P_TYPE] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_SIZE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_SIZE].kind() == Datum::NONE)
+ {
+ std::uniform_int_distribution<int32_t> dist(1, 50);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_SIZE));
+ int32_t *p_size = reinterpret_cast<int32_t *>(
+ tld.part[PART::P_SIZE].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.part_to_generate; i++)
+ p_size[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status P_CONTAINER(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_CONTAINER].kind() == Datum::NONE)
+ {
+ std::uniform_int_distribution<int> dist1(0, static_cast<uint8_t>(kNumContainers_1 - 1));
+ std::uniform_int_distribution<int> dist2(0, static_cast<uint8_t>(kNumContainers_2 - 1));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_CONTAINER));
+ char *p_container = reinterpret_cast<char *>(
+ tld.part[PART::P_CONTAINER].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*part_types_[PART::P_CONTAINER]);
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ int container1_idx = dist1(tld.rng);
+ int container2_idx = dist2(tld.rng);
+ const char *container1 = Containers_1[container1_idx];
+ const char *container2 = Containers_2[container2_idx];
+ size_t container1_length = std::strlen(container1);
+ size_t container2_length = std::strlen(container2);
+
+ char *row = p_container + byte_width * irow;
+ std::strncpy(row, container1, byte_width);
+ std::memcpy(row + container1_length, container2, container2_length);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_RETAILPRICE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_RETAILPRICE].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_RETAILPRICE));
+ const int32_t *p_partkey = reinterpret_cast<const int32_t *>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+ Decimal128 *p_retailprice = reinterpret_cast<Decimal128 *>(
+ tld.part[PART::P_RETAILPRICE].array()->buffers[1]->mutable_data());
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ int32_t partkey = p_partkey[irow];
+ int64_t retail_price = (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ p_retailprice[irow] = { retail_price };
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_COMMENT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_COMMENT].kind() == Datum::NONE)
+ {
+ ARROW_ASSIGN_OR_RAISE(tld.part[PART::P_COMMENT], g_text.GenerateComments(tld.part_to_generate, 5, 22, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ int64_t PartsuppBatchesToGenerate(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t num_batches = (ps_to_generate + batch_size_ - 1) / batch_size_;
+ return num_batches;
+ }
+
+ Status InitPartsupp(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ tld.generated_partsupp.reset();
+ int64_t num_batches = PartsuppBatchesToGenerate(thread_index);
+ tld.partsupp.resize(num_batches);
+ for(std::vector<Datum> &batch : tld.partsupp)
+ {
+ batch.resize(PARTSUPP::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocatePartSuppBatch(size_t thread_index, size_t ibatch, int column)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ int32_t byte_width = arrow::internal::GetByteWidth(*partsupp_types_[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff, AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(partsupp_types_[column], batch_size_, { nullptr, std::move(buff) });
+ tld.partsupp[ibatch][column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status PS_PARTKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_partsupp[PARTSUPP::PS_PARTKEY])
+ {
+ tld.generated_partsupp[PARTSUPP::PS_PARTKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t *p_partkey = reinterpret_cast<const int32_t *>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for(int64_t irow = 0; irow < ps_to_generate; ibatch++)
+ {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_PARTKEY));
+ int32_t *ps_partkey = reinterpret_cast<int32_t *>(
+ tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY].array()->buffers[1]->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; ipartsupp < kPartSuppRowsPerPart && irun < next_run; ipartsupp++, irun++)
+ ps_partkey[batch_offset++] = p_partkey[ipart];
+
+ if(ipartsupp == kPartSuppRowsPerPart)
+ {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_partsupp[PARTSUPP::PS_SUPPKEY])
+ {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t *p_partkey = reinterpret_cast<const int32_t *>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for(int64_t irow = 0; irow < ps_to_generate; ibatch++)
+ {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPKEY));
+ int32_t *ps_suppkey = reinterpret_cast<int32_t *>(
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY].array()->buffers[1]->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; ipartsupp < kPartSuppRowsPerPart && irun < next_run; ipartsupp++, irun++)
+ {
+ int32_t supplier = static_cast<int32_t>(ipartsupp);
+ int32_t partkey = p_partkey[ipart];
+ ps_suppkey[batch_offset++] = (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ if(ipartsupp == kPartSuppRowsPerPart)
+ {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_AVAILQTY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_partsupp[PARTSUPP::PS_AVAILQTY])
+ {
+ tld.generated_partsupp[PARTSUPP::PS_AVAILQTY] = true;
+ std::uniform_int_distribution<int32_t> dist(1, 9999);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for(int64_t irow = 0; irow < ps_to_generate; ibatch++)
+ {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_AVAILQTY));
+ int32_t *ps_availqty = reinterpret_cast<int32_t *>(
+ tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY].array()->buffers[1]->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for(int64_t irun = 0; irun < next_run; irun++)
+ ps_availqty[irun] = dist(tld.rng);
+
+ tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPLYCOST(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST])
+ {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST] = true;
+ std::uniform_int_distribution<int64_t> dist(100, 100000);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for(int64_t irow = 0; irow < ps_to_generate; ibatch++)
+ {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPLYCOST));
+ Decimal128 *ps_supplycost = reinterpret_cast<Decimal128 *>(
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST].array()->buffers[1]->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for(int64_t irun = 0; irun < next_run; irun++)
+ ps_supplycost[irun] = { dist(tld.rng) };
+
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_COMMENT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PARTSUPP::PS_COMMENT].kind() == Datum::NONE)
+ {
+ int64_t irow = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for(size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++)
+ {
+ int64_t num_rows = std::min(batch_size_, ps_to_generate - irow);
+ ARROW_ASSIGN_OR_RAISE(
+ tld.partsupp[ibatch][PARTSUPP::PS_COMMENT], g_text.GenerateComments(num_rows, 49, 198, tld.rng));
+ irow += num_rows;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData
+ {
+ std::vector<Datum> part;
+ std::vector<int8_t> string_indices;
+ int64_t part_to_generate;
+ int64_t partkey_start;
+
+ std::vector<std::vector<Datum>> partsupp;
+ std::bitset<PARTSUPP::kNumCols> generated_partsupp;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex part_output_queue_mutex_;
+ std::mutex partsupp_output_queue_mutex_;
+ std::queue<ExecBatch> part_output_queue_;
+ std::queue<ExecBatch> partsupp_output_queue_;
+ int64_t batch_size_;
+ float scale_factor_;
+ int64_t part_rows_to_generate_;
+ int64_t part_rows_generated_;
+ std::vector<int> part_cols_;
+ std::vector<int> partsupp_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<size_t> part_batches_generated_ = { 0 };
+ std::atomic<size_t> partsupp_batches_generated_ = { 0 };
+ static constexpr int64_t kPartSuppRowsPerPart = 4;
+ };
+
+ class OrdersAndLineItemGenerator
+ {
+ public:
+ Status Init(
+ size_t num_threads,
+ int64_t batch_size,
+ float scale_factor)
+ {
+ if(!inited_)
+ {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ thread_local_data_.resize(num_threads);
+ for(ThreadLocalData &tld : thread_local_data_)
+ {
+ tld.items_per_order.resize(batch_size_);
+ }
+ orders_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000 * 10);
+ }
+ return Status::OK();
+ }
+
+ int64_t orders_batches_generated() const
+ {
+ return orders_batches_generated_.load();
+ }
+
+ int64_t lineitem_batches_generated() const
+ {
+ return lineitem_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetOrdersOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, orders_types_, orders_name_map_, orders_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetLineItemOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, lineitem_types_, lineitem_name_map_, lineitem_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextOrdersBatch()
+ {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if(!orders_output_queue_.empty())
+ {
+ ExecBatch batch = std::move(orders_output_queue_.front());
+ orders_output_queue_.pop();
+ return std::move(batch);
+ }
+ else if(orders_rows_generated_ == orders_rows_to_generate_)
+ {
+ return util::nullopt;
+ }
+ else
+ {
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate = std::min(
+ batch_size_,
+ orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.first_batch_offset = 0;
+ tld.generated_lineitem.reset();
+
+ for(int col : orders_cols_)
+ RETURN_NOT_OK(orders_generators_[col](thread_index));
+ for(int col : lineitem_cols_)
+ RETURN_NOT_OK(lineitem_generators_[col](thread_index));
+
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for(size_t i = 0; i < orders_cols_.size(); i++)
+ {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ if(!lineitem_cols_.empty())
+ {
+ std::vector<ExecBatch> lineitem_results;
+ for(size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++)
+ {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for(size_t icol = 0; icol < lineitem_cols_.size(); icol++)
+ {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(lineitem_output_queue_mutex_);
+ for(ExecBatch &eb : lineitem_results)
+ {
+ lineitem_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(orders_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextLineItemBatch()
+ {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ ExecBatch queued;
+ bool from_queue = false;
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ if(!lineitem_output_queue_.empty())
+ {
+ queued = std::move(lineitem_output_queue_.front());
+ lineitem_output_queue_.pop();
+ from_queue = true;
+ }
+ }
+ tld.first_batch_offset = 0;
+ if(from_queue)
+ {
+ ARROW_DCHECK(queued.length <= batch_size_);
+ tld.first_batch_offset = queued.length;
+ if(queued.length == batch_size_)
+ return std::move(queued);
+ }
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if(orders_rows_generated_ == orders_rows_to_generate_)
+ {
+ if(from_queue)
+ return std::move(queued);
+ return util::nullopt;
+ }
+
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate = std::min(
+ batch_size_,
+ orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1ll);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.generated_lineitem.reset();
+ if(from_queue)
+ {
+ lineitem_batches_generated_.fetch_sub(1);
+ for(size_t i = 0; i < lineitem_cols_.size(); i++)
+ if(tld.lineitem[0][lineitem_cols_[i]].kind() == Datum::NONE)
+ tld.lineitem[0][lineitem_cols_[i]] = std::move(queued[i]);
+ }
+
+ for(int col : orders_cols_)
+ RETURN_NOT_OK(orders_generators_[col](thread_index));
+ for(int col : lineitem_cols_)
+ RETURN_NOT_OK(lineitem_generators_[col](thread_index));
+
+ if(!orders_cols_.empty())
+ {
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for(size_t i = 0; i < orders_cols_.size(); i++)
+ {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch orders_batch, ExecBatch::Make(std::move(orders_result)));
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ orders_output_queue_.emplace(std::move(orders_batch));
+ }
+ }
+ std::vector<ExecBatch> lineitem_results;
+ for(size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++)
+ {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for(size_t icol = 0; icol < lineitem_cols_.size(); icol++)
+ {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ lineitem_batches_generated_.fetch_add(static_cast<int64_t>(lineitem_results.size()));
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ for(size_t i = 1; i < lineitem_results.size(); i++)
+ lineitem_output_queue_.emplace(std::move(lineitem_results[i]));
+ }
+ return std::move(lineitem_results[0]);
+ }
+
+ private:
+#define FOR_EACH_ORDERS_COLUMN(F) \
+ F(O_ORDERKEY) \
+ F(O_CUSTKEY) \
+ F(O_ORDERSTATUS) \
+ F(O_TOTALPRICE) \
+ F(O_ORDERDATE) \
+ F(O_ORDERPRIORITY) \
+ F(O_CLERK) \
+ F(O_SHIPPRIORITY) \
+ F(O_COMMENT)
+
+#define FOR_EACH_LINEITEM_COLUMN(F) \
+ F(L_ORDERKEY) \
+ F(L_PARTKEY) \
+ F(L_SUPPKEY) \
+ F(L_LINENUMBER) \
+ F(L_QUANTITY) \
+ F(L_EXTENDEDPRICE) \
+ F(L_DISCOUNT) \
+ F(L_TAX) \
+ F(L_RETURNFLAG) \
+ F(L_LINESTATUS) \
+ F(L_SHIPDATE) \
+ F(L_COMMITDATE) \
+ F(L_RECEIPTDATE) \
+ F(L_SHIPINSTRUCT) \
+ F(L_SHIPMODE) \
+ F(L_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct ORDERS
+ {
+ enum
+ {
+ FOR_EACH_ORDERS_COLUMN(MAKE_ENUM)
+ kNumCols,
+ };
+ };
+ struct LINEITEM
+ {
+ enum
+ {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_ENUM)
+ kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) \
+ { #col, ORDERS::col },
+ const std::unordered_map<std::string, int> orders_name_map_ =
+ {
+ FOR_EACH_ORDERS_COLUMN(MAKE_STRING_MAP)
+ };
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) \
+ { #col, LINEITEM::col },
+ const std::unordered_map<std::string, int> lineitem_name_map_ =
+ {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_STRING_MAP)
+ };
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> orders_generators_ =
+ {
+ FOR_EACH_ORDERS_COLUMN(MAKE_FN_ARRAY)
+ };
+ std::vector<GenerateColumnFn> lineitem_generators_ =
+ {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_FN_ARRAY)
+ };
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> orders_types_ =
+ {
+ int32(),
+ int32(),
+ fixed_size_binary(1),
+ decimal(12, 2),
+ date32(),
+ fixed_size_binary(15),
+ fixed_size_binary(15),
+ int32(),
+ utf8()
+ };
+
+ const std::vector<std::shared_ptr<DataType>> lineitem_types_ =
+ {
+ int32(),
+ int32(),
+ int32(),
+ int32(),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ fixed_size_binary(1),
+ fixed_size_binary(1),
+ date32(),
+ date32(),
+ date32(),
+ fixed_size_binary(25),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status AllocateOrdersBatch(size_t thread_index, int column)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.orders[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*orders_types_[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff, AllocateBuffer(tld.orders_to_generate * byte_width));
+ ArrayData ad(orders_types_[column], tld.orders_to_generate, { nullptr, std::move(buff) });
+ tld.orders[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status O_ORDERKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_ORDERKEY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERKEY));
+ int32_t *o_orderkey = reinterpret_cast<int32_t *>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.orders_to_generate; i++)
+ {
+ int32_t orderkey_index = tld.orderkey_start + i;
+ int32_t index_of_run = orderkey_index / 8;
+ int32_t index_in_run = orderkey_index % 8;
+ o_orderkey[i] = (index_of_run * 32 + index_in_run + 1);
+ ARROW_DCHECK(1 <= o_orderkey[i] && o_orderkey[i] <= 4 * orders_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CUSTKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_CUSTKEY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CUSTKEY));
+
+ // Spec says it must be a random number between 1 and SF*150000 that is not
+ // divisible by 3. Rather than repeatedly generating numbers until we get to
+ // a non-divisible-by-3 number, we just generate a number between
+ // 0 and SF * 50000 - 1, multiply by 3, and then add either 1 or 2.
+ int32_t sf_50k = static_cast<int32_t>(scale_factor_ * 50000);
+ std::uniform_int_distribution<int32_t> base_dist(0, sf_50k - 1);
+ std::uniform_int_distribution<int32_t> offset_dist(1, 2);
+ int32_t *o_custkey = reinterpret_cast<int32_t *>(
+ tld.orders[ORDERS::O_CUSTKEY].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.orders_to_generate; i++)
+ o_custkey[i] = 3 * base_dist(tld.rng) + offset_dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERSTATUS(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_ORDERSTATUS].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(L_LINESTATUS(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERSTATUS));
+
+ char *o_orderstatus = reinterpret_cast<char *>(
+ tld.orders[ORDERS::O_ORDERSTATUS].array()->buffers[1]->mutable_data());
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ bool all_f = true;
+ bool all_o = true;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ const char *l_linestatus = reinterpret_cast<const char *>(
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->buffers[1]->data());
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++, batch_offset++)
+ {
+ all_f &= l_linestatus[batch_offset] == 'F';
+ all_o &= l_linestatus[batch_offset] == 'O';
+ }
+ if(iline == tld.items_per_order[iorder])
+ {
+ iline = 0;
+ ARROW_DCHECK(!(all_f && all_o));
+ if(all_f)
+ o_orderstatus[iorder] = 'F';
+ else if(all_o)
+ o_orderstatus[iorder] = 'O';
+ else
+ o_orderstatus[iorder] = 'P';
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_TOTALPRICE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_TOTALPRICE].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(L_EXTENDEDPRICE(thread_index));
+ RETURN_NOT_OK(L_TAX(thread_index));
+ RETURN_NOT_OK(L_DISCOUNT(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_TOTALPRICE));
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ int64_t sum = 0;
+ Decimal128 *o_totalprice = reinterpret_cast<Decimal128 *>(
+ tld.orders[ORDERS::O_TOTALPRICE].array()->buffers[1]->mutable_data());
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ const Decimal128 *l_extendedprice = reinterpret_cast<const Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->buffers[1]->data());
+ const Decimal128 *l_tax = reinterpret_cast<const Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->data());
+ const Decimal128 *l_discount = reinterpret_cast<const Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->buffers[1]->data());
+
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++, batch_offset++)
+ {
+ int64_t eprice = static_cast<int64_t>(l_extendedprice[batch_offset]);
+ int64_t tax = static_cast<int64_t>(l_tax[batch_offset]);
+ int64_t discount = static_cast<int64_t>(l_discount[batch_offset]);
+ sum += (eprice * (100 + tax) * (100 - discount));
+ }
+ if(iline == tld.items_per_order[iorder])
+ {
+ sum /= 100 * 100;
+ o_totalprice[iorder] = { sum };
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERDATE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_ORDERDATE].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERDATE));
+
+ std::uniform_int_distribution<uint32_t> dist(STARTDATE, ENDDATE - 151);
+ uint32_t *o_orderdate = reinterpret_cast<uint32_t *>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.orders_to_generate; i++)
+ o_orderdate[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERPRIORITY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_ORDERPRIORITY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERPRIORITY));
+ int32_t byte_width = arrow::internal::GetByteWidth(*orders_types_[ORDERS::O_ORDERPRIORITY]);
+ std::uniform_int_distribution<int32_t> dist(0, kNumPriorities - 1);
+ char *o_orderpriority = reinterpret_cast<char *>(
+ tld.orders[ORDERS::O_ORDERPRIORITY].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.orders_to_generate; i++)
+ {
+ const char *str = Priorities[dist(tld.rng)];
+ std::strncpy(o_orderpriority + i * byte_width, str, byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CLERK(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_CLERK].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CLERK));
+ int32_t byte_width = arrow::internal::GetByteWidth(*orders_types_[ORDERS::O_CLERK]);
+ int64_t max_clerk_id = static_cast<int64_t>(scale_factor_ * 1000);
+ std::uniform_int_distribution<int64_t> dist(1, max_clerk_id);
+ char *o_clerk = reinterpret_cast<char *>(
+ tld.orders[ORDERS::O_CLERK].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.orders_to_generate; i++)
+ {
+ const char *clerk = "Clerk#";
+ const size_t clerk_length = std::strlen(clerk);
+ int64_t clerk_number = dist(tld.rng);
+ char *output = o_clerk + i * byte_width;
+ std::strncpy(output, clerk, byte_width);
+ AppendNumberPaddedToNineDigits(output + clerk_length, clerk_number);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_SHIPPRIORITY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_SHIPPRIORITY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_SHIPPRIORITY));
+ int32_t *o_shippriority = reinterpret_cast<int32_t *>(
+ tld.orders[ORDERS::O_SHIPPRIORITY].array()->buffers[1]->mutable_data());
+ std::memset(o_shippriority, 0, tld.orders_to_generate * sizeof(int32_t));
+ }
+ return Status::OK();
+ }
+
+ Status O_COMMENT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_COMMENT].kind() == Datum::NONE)
+ {
+ ARROW_ASSIGN_OR_RAISE(tld.orders[ORDERS::O_COMMENT], g_text.GenerateComments(tld.orders_to_generate, 19, 78, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ Status GenerateRowCounts(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ std::uniform_int_distribution<int> length_dist(1, 7);
+ tld.lineitem_to_generate = 0;
+ tld.items_per_order.clear();
+ for(int64_t i = 0; i < tld.orders_to_generate; i++)
+ {
+ int64_t length = length_dist(tld.rng);
+ tld.items_per_order.push_back(length);
+ tld.lineitem_to_generate += length;
+ }
+ int64_t num_batches = (tld.first_batch_offset + tld.lineitem_to_generate + batch_size_ - 1) / batch_size_;
+ tld.lineitem.resize(num_batches);
+ for(std::vector<Datum> &batch : tld.lineitem)
+ {
+ batch.resize(LINEITEM::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocateLineItemBufferIfNeeded(size_t thread_index, size_t ibatch, int column, size_t &out_batch_offset)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.lineitem[ibatch][column].kind() == Datum::NONE)
+ {
+ int32_t byte_width = arrow::internal::GetByteWidth(*lineitem_types_[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff, AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(lineitem_types_[column], batch_size_, { nullptr, std::move(buff) });
+ tld.lineitem[ibatch][column] = std::move(ad);
+ out_batch_offset = 0;
+ }
+ if(ibatch == 0)
+ out_batch_offset = tld.first_batch_offset;
+
+ return Status::OK();
+ }
+
+ Status L_ORDERKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_ORDERKEY])
+ {
+ tld.generated_lineitem[LINEITEM::L_ORDERKEY] = true;
+ RETURN_NOT_OK(O_ORDERKEY(thread_index));
+ const int32_t *o_orderkey = reinterpret_cast<const int32_t *>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_ORDERKEY, batch_offset));
+ int32_t *l_linenumber = reinterpret_cast<int32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_ORDERKEY].array()->buffers[1]->mutable_data());
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_linenumber[batch_offset++] = o_orderkey[iorder];
+ if(iline == tld.items_per_order[iorder])
+ {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_ORDERKEY].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_PARTKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_PARTKEY])
+ {
+ tld.generated_lineitem[LINEITEM::L_PARTKEY] = true;
+
+ size_t ibatch = 0;
+ int32_t max_partkey = static_cast<int32_t>(scale_factor_ * 200000);
+ std::uniform_int_distribution<int32_t> dist(1, max_partkey);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_PARTKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t *l_partkey = reinterpret_cast<int32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_partkey[batch_offset] = dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SUPPKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_SUPPKEY])
+ {
+ tld.generated_lineitem[LINEITEM::L_SUPPKEY] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 3);
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset = 0;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_SUPPKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t *l_suppkey = reinterpret_cast<int32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_SUPPKEY].array()->buffers[1]->mutable_data());
+ const int32_t *l_partkey = reinterpret_cast<const int32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ for(int64_t i = 0; i < next_run; i++)
+ {
+ int32_t supplier = dist(tld.rng);
+ int32_t partkey = l_partkey[batch_offset];
+ // Fun fact: the parentheses for this expression are unbalanced in the TPC-H spec.
+ l_suppkey[batch_offset++] = (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SUPPKEY].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINENUMBER(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_LINENUMBER])
+ {
+ tld.generated_lineitem[LINEITEM::L_LINENUMBER] = true;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_LINENUMBER, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t *l_linenumber = reinterpret_cast<int32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_LINENUMBER].array()->buffers[1]->mutable_data());
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ {
+ l_linenumber[batch_offset++] = (iline + 1);
+ ARROW_DCHECK(1 <= (iline + 1) && (iline + 1) <= 7);
+ }
+ if(iline == tld.items_per_order[iorder])
+ {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINENUMBER].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_QUANTITY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_QUANTITY])
+ {
+ tld.generated_lineitem[LINEITEM::L_QUANTITY] = true;
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int64_t> dist(1, 50);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_QUANTITY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128 *l_quantity = reinterpret_cast<Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < next_run; i++)
+ {
+ // Multiply by 100 because the type is decimal(12, 2), so the decimal goes after two digits
+ int64_t quantity = dist(tld.rng) * 100;
+ l_quantity[batch_offset++] = { quantity };
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_EXTENDEDPRICE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE])
+ {
+ tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+ RETURN_NOT_OK(L_QUANTITY(thread_index));
+ size_t ibatch = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_EXTENDEDPRICE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ const int32_t *l_partkey = reinterpret_cast<const int32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ const Decimal128 *l_quantity = reinterpret_cast<const Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->buffers[1]->data());
+ Decimal128 *l_extendedprice = reinterpret_cast<Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ {
+ int64_t partkey = static_cast<int64_t>(l_partkey[batch_offset]);
+ // Divide by 100 to recover the integer representation (not Decimal).
+ int64_t quantity = static_cast<int64_t>(l_quantity[batch_offset]) / 100;
+
+ // Spec says to divide by 100, but that happens automatically due to this being stored
+ // to two decimal points.
+ int64_t retail_price = (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ int64_t extended_price = retail_price * quantity;
+ l_extendedprice[batch_offset] = { extended_price };
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_DISCOUNT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_DISCOUNT])
+ {
+ tld.generated_lineitem[LINEITEM::L_DISCOUNT] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 10);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_DISCOUNT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128 *l_discount = reinterpret_cast<Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_discount[batch_offset] = { dist(tld.rng) };
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_TAX(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_TAX])
+ {
+ tld.generated_lineitem[LINEITEM::L_TAX] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 8);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_TAX, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128 *l_tax = reinterpret_cast<Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_tax[batch_offset] = { dist(tld.rng) };
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RETURNFLAG(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_RETURNFLAG])
+ {
+ tld.generated_lineitem[LINEITEM::L_RETURNFLAG] = true;
+ RETURN_NOT_OK(L_RECEIPTDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<uint32_t> dist;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_RETURNFLAG, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char *l_returnflag = reinterpret_cast<char *>(
+ tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG].array()->buffers[1]->mutable_data());
+ const uint32_t *l_receiptdate = reinterpret_cast<const uint32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ {
+ if(l_receiptdate[batch_offset] <= CURRENTDATE)
+ {
+ uint32_t r = dist(tld.rng);
+ l_returnflag[batch_offset] = (r % 2 == 1) ? 'R' : 'A';
+ }
+ else
+ {
+ l_returnflag[batch_offset] = 'N';
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINESTATUS(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_LINESTATUS])
+ {
+ tld.generated_lineitem[LINEITEM::L_LINESTATUS] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_LINESTATUS, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char *l_linestatus = reinterpret_cast<char *>(
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->buffers[1]->mutable_data());
+ const uint32_t *l_shipdate = reinterpret_cast<const uint32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ {
+ if(l_shipdate[batch_offset] > CURRENTDATE)
+ l_linestatus[batch_offset] = 'O';
+ else
+ l_linestatus[batch_offset] = 'F';
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPDATE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_SHIPDATE])
+ {
+ tld.generated_lineitem[LINEITEM::L_SHIPDATE] = true;
+ RETURN_NOT_OK(O_ORDERDATE(thread_index));
+ const int32_t *o_orderdate = reinterpret_cast<const int32_t *>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(1, 121);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_SHIPDATE, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t *l_shipdate = reinterpret_cast<uint32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->buffers[1]->mutable_data());
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_shipdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if(iline == tld.items_per_order[iorder])
+ {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMITDATE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_COMMITDATE])
+ {
+ tld.generated_lineitem[LINEITEM::L_COMMITDATE] = true;
+ const int32_t *o_orderdate = reinterpret_cast<const int32_t *>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(30, 90);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_COMMITDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t *l_commitdate = reinterpret_cast<uint32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_COMMITDATE].array()->buffers[1]->mutable_data());
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_commitdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if(iline == tld.items_per_order[iorder])
+ {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMITDATE].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RECEIPTDATE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_RECEIPTDATE])
+ {
+ tld.generated_lineitem[LINEITEM::L_RECEIPTDATE] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(1, 30);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_RECEIPTDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t *l_receiptdate = reinterpret_cast<uint32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE].array()->buffers[1]->mutable_data());
+ const uint32_t *l_shipdate = reinterpret_cast<const uint32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_receiptdate[batch_offset] = l_shipdate[batch_offset] + dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPINSTRUCT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT])
+ {
+ tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT] = true;
+ int32_t byte_width = arrow::internal::GetByteWidth(*lineitem_types_[LINEITEM::L_SHIPINSTRUCT]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumInstructions - 1);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_SHIPINSTRUCT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char *l_shipinstruct = reinterpret_cast<char *>(
+ tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ {
+ const char *str = Instructions[dist(tld.rng)];
+ // Note that we don't have to memset the buffer to 0 because strncpy pads each string
+ // with 0's anyway
+ std::strncpy(l_shipinstruct + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPMODE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_SHIPMODE])
+ {
+ tld.generated_lineitem[LINEITEM::L_SHIPMODE] = true;
+ int32_t byte_width = arrow::internal::GetByteWidth(*lineitem_types_[LINEITEM::L_SHIPMODE]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumModes - 1);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_SHIPMODE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char *l_shipmode = reinterpret_cast<char *>(
+ tld.lineitem[ibatch][LINEITEM::L_SHIPMODE].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ {
+ const char *str = Modes[dist(tld.rng)];
+ std::strncpy(l_shipmode + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPMODE].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMENT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_COMMENT])
+ {
+ tld.generated_lineitem[LINEITEM::L_COMMENT] = true;
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ // Comments are kind of sneaky: we always generate the full batch and then just bump the length
+ if(tld.lineitem[ibatch][LINEITEM::L_COMMENT].kind() == Datum::NONE)
+ {
+ ARROW_ASSIGN_OR_RAISE(tld.lineitem[ibatch][LINEITEM::L_COMMENT], g_text.GenerateComments(batch_size_, 10, 43, tld.rng));
+ batch_offset = 0;
+ }
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ batch_offset += next_run;
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMENT].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData
+ {
+ std::vector<Datum> orders;
+ int64_t orders_to_generate;
+ int64_t orderkey_start;
+
+ std::vector<std::vector<Datum>> lineitem;
+ std::vector<int> items_per_order;
+ int64_t lineitem_to_generate;
+ int64_t first_batch_offset;
+ std::bitset<LINEITEM::kNumCols> generated_lineitem;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex orders_output_queue_mutex_;
+ std::mutex lineitem_output_queue_mutex_;
+ std::queue<ExecBatch> orders_output_queue_;
+ std::queue<ExecBatch> lineitem_output_queue_;
+ int64_t batch_size_;
+ float scale_factor_;
+ int64_t orders_rows_to_generate_;
+ int64_t orders_rows_generated_;
+ std::vector<int> orders_cols_;
+ std::vector<int> lineitem_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<size_t> orders_batches_generated_ = { 0 };
+ std::atomic<size_t> lineitem_batches_generated_ = { 0 };
+ };
+
+ class SupplierGenerator : public TpchTableGenerator
+ {
+ public:
+ Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) override
+ {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 10000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_, SetOutputColumns(
+ columns,
+ types_,
+ name_map_,
+ gen_list_));
+
+ random::pcg32_fast rng;
+ std::uniform_int_distribution<int64_t> dist(0, rows_to_generate_ - 1);
+ size_t num_special_rows = static_cast<size_t>(5 * scale_factor_);
+ std::unordered_set<int64_t> good_rows_set;
+ while(good_rows_set.size() < num_special_rows)
+ {
+ int64_t row = dist(rng);
+ good_rows_set.insert(row);
+ }
+ std::unordered_set<int64_t> bad_rows_set;
+ while(bad_rows_set.size() < num_special_rows)
+ {
+ int64_t bad_row;
+ do
+ {
+ bad_row = dist(rng);
+ } while(good_rows_set.find(bad_row) != good_rows_set.end());
+ bad_rows_set.insert(bad_row);
+ }
+ good_rows_.clear();
+ bad_rows_.clear();
+ good_rows_.insert(good_rows_.end(), good_rows_set.begin(), good_rows_set.end());
+ bad_rows_.insert(bad_rows_.end(), bad_rows_set.begin(), bad_rows_set.end());
+ std::sort(good_rows_.begin(), good_rows_.end());
+ std::sort(bad_rows_.begin(), bad_rows_.end());
+ return Status::OK();
+ }
+
+ Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override
+ {
+ thread_local_data_.resize(num_threads);
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for(size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override
+ {
+ return schema_;
+ }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(S_SUPPKEY) \
+ F(S_NAME) \
+ F(S_ADDRESS) \
+ F(S_NATIONKEY) \
+ F(S_PHONE) \
+ F(S_ACCTBAL) \
+ F(S_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct SUPPLIER
+ {
+ enum
+ {
+ FOR_EACH_COLUMN(MAKE_ENUM)
+ kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) \
+ { #col, SUPPLIER::col },
+ const std::unordered_map<std::string, int> name_map_ =
+ {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)
+ };
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> generators_ =
+ {
+ FOR_EACH_COLUMN(MAKE_FN_ARRAY)
+ };
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> types_ =
+ {
+ int32(),
+ fixed_size_binary(25),
+ utf8(),
+ int32(),
+ fixed_size_binary(15),
+ decimal(12, 2),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index)
+ {
+ if(done_.load())
+ return Status::OK();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ tld.suppkey_start = rows_generated_.fetch_add(batch_size_);
+ if(tld.suppkey_start >= rows_to_generate_)
+ return Status::OK();
+
+ tld.to_generate = std::min(batch_size_,
+ rows_to_generate_ - tld.suppkey_start);
+
+ tld.batch.resize(SUPPLIER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for(int col : gen_list_)
+ RETURN_NOT_OK(generators_[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for(size_t i = 0; i < gen_list_.size(); i++)
+ {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_outputted_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_outputted_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if(is_last_batch)
+ {
+ done_.store(true);
+ finished_callback_(batches_outputted_.load());
+ return Status::OK();
+ }
+ return schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*types_[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff, AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(types_[column], tld.to_generate, { nullptr, std::move(buff) });
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status S_SUPPKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[SUPPLIER::S_SUPPKEY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_SUPPKEY));
+ int32_t *s_suppkey = reinterpret_cast<int32_t *>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->mutable_data());
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ {
+ s_suppkey[irow] = (tld.suppkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_NAME(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[SUPPLIER::S_NAME].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(S_SUPPKEY(thread_index));
+ const int32_t *s_suppkey = reinterpret_cast<const int32_t *>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->data());
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NAME));
+ int32_t byte_width = arrow::internal::GetByteWidth(*types_[SUPPLIER::S_NAME]);
+ char *s_name = reinterpret_cast<char *>(
+ tld.batch[SUPPLIER::S_NAME].array()->buffers[1]->mutable_data());
+ // Look man, I'm just following the spec ok? Section 4.2.3 as of March 1 2022
+ const char *supplier = "Supplie#r";
+ const size_t supplier_length = std::strlen(supplier);
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ {
+ char *out = s_name + byte_width * irow;
+ std::strncpy(out, supplier, byte_width);
+ AppendNumberPaddedToNineDigits(out + supplier_length, s_suppkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ADDRESS(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[SUPPLIER::S_ADDRESS].kind() == Datum::NONE)
+ {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.batch[SUPPLIER::S_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status S_NATIONKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[SUPPLIER::S_NATIONKEY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t *s_nationkey = reinterpret_cast<int32_t *>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->mutable_data());
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status S_PHONE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[SUPPLIER::S_PHONE].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(S_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*types_[SUPPLIER::S_PHONE]);
+ const int32_t *s_nationkey = reinterpret_cast<const int32_t *>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->data());
+ char *s_phone = reinterpret_cast<char *>(
+ tld.batch[SUPPLIER::S_PHONE].array()->buffers[1]->mutable_data());
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ {
+ GeneratePhoneNumber(
+ s_phone + irow * byte_width,
+ tld.rng,
+ s_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ACCTBAL(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[SUPPLIER::S_ACCTBAL].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_ACCTBAL));
+ Decimal128 *s_acctbal = reinterpret_cast<Decimal128 *>(
+ tld.batch[SUPPLIER::S_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_acctbal[irow] = { dist(tld.rng) };
+ }
+ return Status::OK();
+ }
+
+ Status S_COMMENT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[SUPPLIER::S_COMMENT].kind() == Datum::NONE)
+ {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[SUPPLIER::S_COMMENT], g_text.GenerateComments(tld.to_generate, 25, 100, tld.rng));
+ ModifyComments(thread_index, "Recommends", good_rows_);
+ ModifyComments(thread_index, "Complaints", bad_rows_);
+ }
+ return Status::OK();
+ }
+
+ void ModifyComments(
+ size_t thread_index,
+ const char *review,
+ const std::vector<int64_t> &indices)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ const int32_t *offsets = reinterpret_cast<const int32_t *>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[1]->data());
+ char *str = reinterpret_cast<char *>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[2]->mutable_data());
+ const char *customer = "Customer";
+ const size_t customer_length = std::strlen(customer);
+ const size_t review_length = std::strlen(review);
+
+ auto it = std::lower_bound(indices.begin(), indices.end(), tld.suppkey_start);
+ for(; it != indices.end() && *it < tld.suppkey_start + tld.to_generate; it++)
+ {
+ int64_t idx_in_batch = *it - tld.suppkey_start;
+ char *out = str + offsets[idx_in_batch];
+ int32_t str_length = offsets[idx_in_batch + 1] - offsets[idx_in_batch];
+ std::uniform_int_distribution<int32_t> gap_dist(0, str_length - customer_length - review_length);
+ int32_t gap = gap_dist(tld.rng);
+ int32_t total_length = customer_length + gap + review_length;
+ std::uniform_int_distribution<int32_t> start_dist(0, str_length - total_length);
+ int32_t start = start_dist(tld.rng);
+ std::memcpy(out + start, customer, customer_length);
+ std::memcpy(out + start + customer_length + gap, review, review_length);
+ }
+ }
+
+ struct ThreadLocalData
+ {
+ random::pcg32_fast rng;
+ int64_t suppkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+ std::vector<int64_t> good_rows_;
+ std::vector<int64_t> bad_rows_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_;
+ std::atomic<int64_t> rows_generated_;
+ float scale_factor_;
+ int64_t batch_size_;
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+ };
+
+ class PartGenerator : public TpchTableGenerator
+ {
+ public:
+ PartGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen))
+ {
+ }
+
+ Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) override
+ {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ gen_->SetPartOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override
+ {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for(size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override
+ {
+ return schema_;
+ }
+
+ private:
+ Status ProduceCallback(size_t)
+ {
+ if(done_.load())
+ return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextPartBatch());
+ if(!maybe_batch.has_value())
+ {
+ int64_t batches_generated = gen_->part_batches_generated();
+ if(batches_generated == batches_outputted_.load())
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+ };
+
+ class PartSuppGenerator : public TpchTableGenerator
+ {
+ public:
+ PartSuppGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen))
+ {
+ }
+
+ Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) override
+ {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ gen_->SetPartSuppOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override
+ {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for(size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override
+ {
+ return schema_;
+ }
+
+ private:
+ Status ProduceCallback(size_t)
+ {
+ if(done_.load())
+ return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextPartSuppBatch());
+ if(!maybe_batch.has_value())
+ {
+ int64_t batches_generated = gen_->partsupp_batches_generated();
+ if(batches_generated == batches_outputted_.load())
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+ };
+
+ class CustomerGenerator : public TpchTableGenerator
+ {
+ public:
+ Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) override
+ {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = scale_factor_ * 150000;
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_, SetOutputColumns(
+ columns,
+ types_,
+ name_map_,
+ gen_list_));
+ return Status::OK();
+ }
+
+ Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override
+ {
+ thread_local_data_.resize(num_threads);
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for(size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override
+ {
+ return schema_;
+ }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(C_CUSTKEY) \
+ F(C_NAME) \
+ F(C_ADDRESS) \
+ F(C_NATIONKEY) \
+ F(C_PHONE) \
+ F(C_ACCTBAL) \
+ F(C_MKTSEGMENT) \
+ F(C_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct CUSTOMER
+ {
+ enum
+ {
+ FOR_EACH_COLUMN(MAKE_ENUM)
+ kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) \
+ { #col, CUSTOMER::col },
+ const std::unordered_map<std::string, int> name_map_ =
+ {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)
+ };
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> generators_ =
+ {
+ FOR_EACH_COLUMN(MAKE_FN_ARRAY)
+ };
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> types_ =
+ {
+ int32(),
+ utf8(),
+ utf8(),
+ int32(),
+ fixed_size_binary(15),
+ decimal(12, 2),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index)
+ {
+ if(done_.load())
+ return Status::OK();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ tld.custkey_start = rows_generated_.fetch_add(batch_size_);
+ if(tld.custkey_start >= rows_to_generate_)
+ return Status::OK();
+
+ tld.to_generate = std::min(batch_size_,
+ rows_to_generate_ - tld.custkey_start);
+
+ tld.batch.resize(CUSTOMER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for(int col : gen_list_)
+ RETURN_NOT_OK(generators_[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for(size_t i = 0; i < gen_list_.size(); i++)
+ {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_generated_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_generated_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if(is_last_batch)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ return schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*types_[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff, AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(types_[column], tld.to_generate, { nullptr, std::move(buff) });
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status C_CUSTKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[CUSTOMER::C_CUSTKEY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_CUSTKEY));
+ int32_t *c_custkey = reinterpret_cast<int32_t *>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->mutable_data());
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ {
+ c_custkey[irow] = (tld.custkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_NAME(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[CUSTOMER::C_NAME].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(C_CUSTKEY(thread_index));
+ const int32_t *c_custkey = reinterpret_cast<const int32_t *>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->data());
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff, AllocateBuffer((tld.to_generate + 1) * sizeof(int32_t)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buff->mutable_data());
+ const char *customer = "Customer#";
+ const size_t customer_length = std::strlen(customer);
+ offsets[0] = 0;
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ {
+ int num_digits = GetNumDigits(c_custkey[irow]);
+ int num_chars = std::max(num_digits, 9);
+ offsets[irow + 1] = offsets[irow] + num_chars + customer_length;
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff, AllocateBuffer(offsets[tld.to_generate]));
+ char *str = reinterpret_cast<char *>(str_buff->mutable_data());
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ {
+ char *out = str + offsets[irow];
+ std::memcpy(out, customer, customer_length);
+ AppendNumberPaddedToNineDigits(out + customer_length, c_custkey[irow]);
+ }
+ ArrayData ad(utf8(), tld.to_generate, { nullptr, std::move(offset_buff), std::move(str_buff) });
+ tld.batch[CUSTOMER::C_NAME] = std::move(ad);
+ }
+ return Status::OK();
+ }
+
+ Status C_ADDRESS(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[CUSTOMER::C_ADDRESS].kind() == Datum::NONE)
+ {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.batch[CUSTOMER::C_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status C_NATIONKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[CUSTOMER::C_NATIONKEY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t *c_nationkey = reinterpret_cast<int32_t *>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->mutable_data());
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status C_PHONE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[CUSTOMER::C_PHONE].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(C_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*types_[CUSTOMER::C_PHONE]);
+ const int32_t *c_nationkey = reinterpret_cast<const int32_t *>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->data());
+ char *c_phone = reinterpret_cast<char *>(
+ tld.batch[CUSTOMER::C_PHONE].array()->buffers[1]->mutable_data());
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ {
+ GeneratePhoneNumber(
+ c_phone + irow * byte_width,
+ tld.rng,
+ c_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_ACCTBAL(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[CUSTOMER::C_ACCTBAL].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_ACCTBAL));
+ Decimal128 *c_acctbal = reinterpret_cast<Decimal128 *>(
+ tld.batch[CUSTOMER::C_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_acctbal[irow] = { dist(tld.rng) };
+ }
+ return Status::OK();
+ }
+
+ Status C_MKTSEGMENT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[CUSTOMER::C_MKTSEGMENT].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_MKTSEGMENT));
+ int32_t byte_width = arrow::internal::GetByteWidth(*types_[CUSTOMER::C_MKTSEGMENT]);
+ char *c_mktsegment = reinterpret_cast<char *>(
+ tld.batch[CUSTOMER::C_MKTSEGMENT].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int32_t> dist(0, kNumSegments - 1);
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ {
+ char *out = c_mktsegment + irow * byte_width;
+ int str_idx = dist(tld.rng);
+ std::strncpy(out, Segments[str_idx], byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_COMMENT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[CUSTOMER::C_COMMENT].kind() == Datum::NONE)
+ {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[CUSTOMER::C_COMMENT], g_text.GenerateComments(tld.to_generate, 29, 116, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData
+ {
+ random::pcg32_fast rng;
+ int64_t custkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_;
+ std::atomic<int64_t> rows_generated_;
+ float scale_factor_;
+ int64_t batch_size_;
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+ };
+
+ class OrdersGenerator : public TpchTableGenerator
+ {
+ public:
+ OrdersGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen))
+ {
+ }
+
+ Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) override
+ {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ gen_->SetOrdersOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override
+ {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for(size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override
+ {
+ return schema_;
+ }
+
+ private:
+ Status ProduceCallback(size_t)
+ {
+ if(done_.load())
+ return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextOrdersBatch());
+ if(!maybe_batch.has_value())
+ {
+ int64_t batches_generated = gen_->orders_batches_generated();
+ if(batches_generated == batches_outputted_.load())
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+ };
+
+ class LineitemGenerator : public TpchTableGenerator
+ {
+ public:
+ LineitemGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen))
+ {}
+
+ Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) override
+ {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ gen_->SetLineItemOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override
+ {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for(size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override
+ {
+ return schema_;
+ }
+
+ private:
+ Status ProduceCallback(size_t)
+ {
+ if(done_.load())
+ return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextLineItemBatch());
+ if(!maybe_batch.has_value())
+ {
+ int64_t batches_generated = gen_->lineitem_batches_generated();
+ if(batches_generated == batches_outputted_.load())
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+ };
+
+ class NationGenerator : public TpchTableGenerator
+ {
+ public:
+ Status Init(
+ std::vector<std::string> columns,
+ float /*scale_factor*/,
+ int64_t /*batch_size*/) override
+ {
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(
+ columns,
+ types_,
+ name_map_,
+ column_indices_));
+ return Status::OK();
+ }
+
+ Status StartProducing(
+ size_t /*num_threads*/,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback /*schedule_task_callback*/) override
+ {
+ std::shared_ptr<Buffer> N_NATIONKEY_buffer = Buffer::Wrap(N_NATIONKEY, sizeof(N_NATIONKEY));
+ ArrayData N_NATIONKEY_arraydata(int32(), kRowCount, { nullptr, std::move(N_NATIONKEY_buffer) });
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> N_NAME_buffer, AllocateBuffer(kRowCount * kNameByteWidth));
+ char *N_NAME = reinterpret_cast<char *>(N_NAME_buffer->mutable_data());
+ for(size_t i = 0; i < kRowCount; i++)
+ std::strncpy(N_NAME + kNameByteWidth * i, country_names_[i], kNameByteWidth);
+ ArrayData N_NAME_arraydata(fixed_size_binary(kNameByteWidth), kRowCount, { nullptr, std::move(N_NAME_buffer) });
+
+ std::shared_ptr<Buffer> N_REGIONKEY_buffer = Buffer::Wrap(N_REGIONKEY, sizeof(N_REGIONKEY));
+ ArrayData N_REGIONKEY_arraydata(int32(), kRowCount, { nullptr, std::move(N_REGIONKEY_buffer) });
+
+ ARROW_ASSIGN_OR_RAISE(Datum N_COMMENT_datum, g_text.GenerateComments(kRowCount, 31, 114, rng_));
+
+ std::vector<Datum> fields =
+ {
+ std::move(N_NATIONKEY_arraydata),
+ std::move(N_NAME_arraydata),
+ std::move(N_REGIONKEY_arraydata),
+ std::move(N_COMMENT_datum)
+ };
+
+ std::vector<Datum> result;
+ for(const int &col : column_indices_)
+ result.push_back(fields[col]);
+ ARROW_ASSIGN_OR_RAISE(ExecBatch batch, ExecBatch::Make(std::move(result)));
+ output_callback(std::move(batch));
+ finished_callback(static_cast<int64_t>(1));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override
+ {
+ return schema_;
+ }
+
+ private:
+ random::pcg32_fast rng_;
+
+ static constexpr size_t kRowCount = 25;
+ static constexpr int32_t kNameByteWidth = 25;
+ const int32_t N_NATIONKEY[kRowCount] = { 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24 };
+ const char *country_names_[kRowCount] =
+ {
+ "ALGERIA", "ARGENTINA", "BRAZIL",
+ "CANADA", "EGYPT", "ETHIOPIA",
+ "FRANCE", "GERMANY", "INDIA",
+ "INDONESIA", "IRAN", "IRAQ",
+ "JAPAN", "JORDAN", "KENYA",
+ "MOROCCO", "MOZAMBIQUE", "PERU",
+ "CHINA", "ROMANIA", "SAUDI ARABIA",
+ "VIETNAM", "RUSSIA", "UNITED KINGDOM",
+ "UNITED STATES"
+ };
+ const int32_t N_REGIONKEY[kRowCount] = { 0, 1, 1, 1, 4, 0, 3, 3, 2, 2, 4, 4, 2, 4, 0, 0, 0, 1, 2, 3, 4, 2, 3, 3, 1 };
+
+ struct NATION
+ {
+ enum
+ {
+ N_NATIONKEY,
+ N_NAME,
+ N_REGIONKEY,
+ N_COMMENT,
+ };
+ };
+
+ const std::unordered_map<std::string, int> name_map_ =
+ {
+ { "N_NATIONKEY", NATION::N_NATIONKEY },
+ { "N_NAME", NATION::N_NAME },
+ { "N_REGIONKEY", NATION::N_REGIONKEY },
+ { "N_COMMENT", NATION::N_COMMENT },
+ };
+
+ std::vector<std::shared_ptr<DataType>> types_ =
+ {
+ int32(),
+ fixed_size_binary(kNameByteWidth),
+ int32(),
+ utf8(),
+ };
+
+ std::shared_ptr<Schema> schema_;
+ std::vector<int> column_indices_;
+ };
+
+ class RegionGenerator : public TpchTableGenerator
+ {
+ public:
+ Status Init(
+ std::vector<std::string> columns,
+ float /*scale_factor*/,
+ int64_t /*batch_size*/) override
+ {
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(
+ columns,
+ types_,
+ name_map_,
+ column_indices_));
+ return Status::OK();
+ }
+
+ Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback /*schedule_task_callback*/) override
+ {
+ std::shared_ptr<Buffer> R_REGIONKEY_buffer = Buffer::Wrap(R_REGIONKEY, sizeof(R_REGIONKEY));
+ ArrayData R_REGIONKEY_arraydata(int32(), kRowCount, { nullptr, std::move(R_REGIONKEY_buffer) });
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> R_NAME_buffer, AllocateBuffer(kRowCount * kNameByteWidth));
+ char *R_NAME_data = reinterpret_cast<char *>(R_NAME_buffer->mutable_data());
+ for(size_t i = 0; i < kRowCount; i++)
+ std::strncpy(R_NAME_data + kNameByteWidth * i, region_names_[i], kNameByteWidth);
+ ArrayData R_NAME_arraydata(types_[static_cast<int>(REGION::R_NAME)], kRowCount, { nullptr, std::move(R_NAME_buffer) });
+
+ ARROW_ASSIGN_OR_RAISE(Datum R_COMMENT_datum, g_text.GenerateComments(kRowCount, 31, 115, rng_));
+
+ std::vector<Datum> fields = { std::move(R_REGIONKEY_arraydata), std::move(R_NAME_arraydata), std::move(R_COMMENT_datum) };
+ std::vector<Datum> result;
+ for(const int &col : column_indices_)
+ result.push_back(fields[col]);
+ ARROW_ASSIGN_OR_RAISE(ExecBatch batch, ExecBatch::Make(std::move(result)));
+ output_callback(std::move(batch));
+ finished_callback(static_cast<int64_t>(1));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override
+ {
+ return schema_;
+ }
+
+ random::pcg32_fast rng_;
+
+ static constexpr size_t kRowCount = 5;
+ static constexpr int32_t kNameByteWidth = 25;
+ const int32_t R_REGIONKEY[kRowCount] = { 0, 1, 2, 3, 4 };
+ const char *region_names_[kRowCount] =
+ {
+ "AFRICA", "AMERICA", "ASIA", "EUROPE", "MIDDLE EAST"
+ };
+
+ struct REGION
+ {
+ enum
+ {
+ R_REGIONKEY,
+ R_NAME,
+ R_COMMENT,
+ kNumColumns,
+ };
+ };
+
+ const std::unordered_map<std::string, int> name_map_ =
+ {
+ { "R_REGIONKEY", REGION::R_REGIONKEY },
+ { "R_NAME", REGION::R_NAME },
+ { "R_COMMENT", REGION::R_COMMENT },
+ };
+
+ const std::vector<std::shared_ptr<DataType>> types_ =
+ {
+ int32(),
+ fixed_size_binary(kNameByteWidth),
+ utf8(),
+ };
+
+ std::shared_ptr<Schema> schema_;
+ std::vector<int> column_indices_;
+ };
+
+ class TpchNode : public ExecNode
+ {
+ public:
+ TpchNode(ExecPlan *plan,
+ std::unique_ptr<TpchTableGenerator> generator)
+ : ExecNode(plan, {}, {}, generator->schema(), /*num_outputs=*/1),
+ generator_(std::move(generator))
+ {
+ }
+
+ const char *kind_name() const override
+ {
+ return "TpchNode";
+ }
+
+ [[noreturn]]
+ static void NoInputs()
+ {
+ Unreachable("TPC-H node should never have any inputs");
+ }
+
+ [[noreturn]]
+ void InputReceived(ExecNode *, ExecBatch) override
+ {
+ NoInputs();
+ }
+
+ [[noreturn]]
+ void ErrorReceived(ExecNode *, Status) override
+ {
+ NoInputs();
+ }
+
+ [[noreturn]]
+ void InputFinished(ExecNode *, int) override
+ {
+ NoInputs();
+ }
+
+ Status StartProducing() override
+ {
+ finished_ = Future<>::Make();
+ return generator_->StartProducing(
+ thread_indexer_.Capacity(),
+ [this](ExecBatch batch) { this->OutputBatchCallback(std::move(batch)); },
+ [this](int64_t num_batches) { this->FinishedCallback(num_batches); },
+ [this](std::function<Status(size_t)> func) -> Status { return this->ScheduleTaskCallback(std::move(func)); }
+ );
+ }
+
+ void PauseProducing(ExecNode *output) override {}
+ void ResumeProducing(ExecNode *output) override {}
+
+ void StopProducing(ExecNode *output) override
+ {
+ DCHECK_EQ(output, outputs_[0]);
+ StopProducing();
+ }
+
+ void StopProducing() override
+ {
+ generator_->Abort([this]() { this->finished_.MarkFinished(); });
+ }
+
+ Future<> finished() override
+ {
+ return finished_;
+ }
+
+ private:
+ void OutputBatchCallback(ExecBatch batch)
+ {
+ outputs_[0]->InputReceived(this, std::move(batch));
+ }
+
+ void FinishedCallback(int64_t total_num_batches)
+ {
+ outputs_[0]->InputFinished(this, static_cast<int>(total_num_batches));
+ finished_.MarkFinished();
+ }
+
+ Status ScheduleTaskCallback(std::function<Status(size_t)> func)
+ {
+ auto executor = plan_->exec_context()->executor();
+ if (executor)
+ {
+ RETURN_NOT_OK(executor->Spawn([this, func]
+ {
+ size_t thread_index = thread_indexer_();
+ Status status = func(thread_index);
+ if (!status.ok())
+ {
+ StopProducing();
+ ErrorIfNotOk(status);
+ return;
+ }
+ }));
+ }
+ else
+ {
+ return func(0);
+ }
+ return Status::OK();
+ }
Review comment:
Ok - this should be done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1070989318
> Thanks @save-buffer I agree, that will be simpler. I'm happy to take a look at tpch_dbgen_write too.
Also agree we should move this forward if the R stuff is holding it back. And would love if you could take a look at that Weston!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829319799
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
Review comment:
none of the other benchmarks use it, so i just copied that. I can try it though
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829337783
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
Review comment:
from google style guide:
> Parameters are either inputs to the function, outputs from the function, or both. Non-optional input parameters should usually be values or const references, while non-optional output and input/output parameters should usually be references (which cannot be null). Generally, use std::optional to represent optional by-value inputs, and use a const pointer when the non-optional form would have used a reference. Use non-const pointers to represent optional outputs and optional input/output parameters.
this is a non-optional input/output parameter, so according to them it should be a reference.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] nealrichardson commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
nealrichardson commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1064081821
Presumably this involves vendoring some official TPC code? If so, we should confirm its license and add a note to LICENSE.txt
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1077334837
I get the following warning when compiling the benchmarks:
```
[241/415] Building CXX object src/arrow/util/CMakeFiles/arrow-small-vector-benchmark.dir/small_vector_benchmark.cc.o
/home/antoine/arrow/dev/cpp/src/arrow/util/small_vector_benchmark.cc:70:37: warning: address of stack memory associated with parameter 'v' returned [-Wreturn-stack-address]
return reinterpret_cast<intptr_t>(v.data());
^
/home/antoine/arrow/dev/cpp/src/arrow/util/small_vector_benchmark.cc:119:16: note: in instantiation of function template specialization 'arrow::internal::ConsumeVector<std::vector<int>>' requested here
dummy += ConsumeVector(vec);
^
/home/antoine/arrow/dev/cpp/src/arrow/util/small_vector_benchmark.cc:325:1: note: in instantiation of function template specialization 'arrow::internal::CopyEmptyVector<std::vector<int>>' requested here
SHORT_VECTOR_BENCHMARKS(STD_VECTOR);
^
/home/antoine/arrow/dev/cpp/src/arrow/util/small_vector_benchmark.cc:287:22: note: expanded from macro 'SHORT_VECTOR_BENCHMARKS'
BENCHMARK_TEMPLATE(CopyEmptyVector, VEC_TYPE_FACTORY(int)); \
^
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] ursabot edited a comment on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
ursabot edited a comment on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082832095
Benchmark runs are scheduled for baseline = 0deefd1838bf4706232b90b72da010698010f7b6 and contender = 50fab73da58c9f76fbe675d661e56087e2f5fb52. 50fab73da58c9f76fbe675d661e56087e2f5fb52 is a master commit associated with this PR. Results will be available as each benchmark for each run completes.
Conbench compare runs links:
[Finished :arrow_down:0.0% :arrow_up:0.0%] [ec2-t3-xlarge-us-east-2](https://conbench.ursa.dev/compare/runs/3f8ac59fddef426cac31b6b2826ce5a6...e5b78c70b6fe4a699c4cb8b967fe5736/)
[Scheduled] [test-mac-arm](https://conbench.ursa.dev/compare/runs/02b74ab900d440f8b5361300f15ad15c...b870e78b617e41e98c2ade7241e0aa06/)
[Scheduled] [ursa-i9-9960x](https://conbench.ursa.dev/compare/runs/1f7740546747452aa61fece1295aa0e3...f66d8edd43724d9988628cbe4a4e4d50/)
[Scheduled] [ursa-thinkcentre-m75q](https://conbench.ursa.dev/compare/runs/c3c55d6999fe4d7eabec73cbdfa1fad6...4b9f45f6442447d3a9d758598d0e14d3/)
Buildkite builds:
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/415| `50fab73d` ec2-t3-xlarge-us-east-2>
[Scheduled] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-test-mac-arm/builds/401| `50fab73d` test-mac-arm>
[Scheduled] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-thinkcentre-m75q/builds/411| `50fab73d` ursa-thinkcentre-m75q>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/414| `0deefd18` ec2-t3-xlarge-us-east-2>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-test-mac-arm/builds/400| `0deefd18` test-mac-arm>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-thinkcentre-m75q/builds/410| `0deefd18` ursa-thinkcentre-m75q>
Supported benchmarks:
ec2-t3-xlarge-us-east-2: Supported benchmark langs: Python, R. Runs only benchmarks with cloud = True
test-mac-arm: Supported benchmark langs: C++, Python, R
ursa-i9-9960x: Supported benchmark langs: Python, R, JavaScript
ursa-thinkcentre-m75q: Supported benchmark langs: C++, Java
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] ursabot edited a comment on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
ursabot edited a comment on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082832095
Benchmark runs are scheduled for baseline = 0deefd1838bf4706232b90b72da010698010f7b6 and contender = 50fab73da58c9f76fbe675d661e56087e2f5fb52. 50fab73da58c9f76fbe675d661e56087e2f5fb52 is a master commit associated with this PR. Results will be available as each benchmark for each run completes.
Conbench compare runs links:
[Failed] [ec2-t3-xlarge-us-east-2](https://conbench.ursa.dev/compare/runs/3f8ac59fddef426cac31b6b2826ce5a6...e5b78c70b6fe4a699c4cb8b967fe5736/)
[Scheduled] [test-mac-arm](https://conbench.ursa.dev/compare/runs/02b74ab900d440f8b5361300f15ad15c...b870e78b617e41e98c2ade7241e0aa06/)
[Scheduled] [ursa-i9-9960x](https://conbench.ursa.dev/compare/runs/1f7740546747452aa61fece1295aa0e3...f66d8edd43724d9988628cbe4a4e4d50/)
[Scheduled] [ursa-thinkcentre-m75q](https://conbench.ursa.dev/compare/runs/c3c55d6999fe4d7eabec73cbdfa1fad6...4b9f45f6442447d3a9d758598d0e14d3/)
Supported benchmarks:
ec2-t3-xlarge-us-east-2: Supported benchmark langs: Python, R. Runs only benchmarks with cloud = True
test-mac-arm: Supported benchmark langs: C++, Python, R
ursa-i9-9960x: Supported benchmark langs: Python, R, JavaScript
ursa-thinkcentre-m75q: Supported benchmark langs: C++, Java
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] ursabot edited a comment on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
ursabot edited a comment on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082832095
Benchmark runs are scheduled for baseline = 0deefd1838bf4706232b90b72da010698010f7b6 and contender = 50fab73da58c9f76fbe675d661e56087e2f5fb52. 50fab73da58c9f76fbe675d661e56087e2f5fb52 is a master commit associated with this PR. Results will be available as each benchmark for each run completes.
Conbench compare runs links:
[Finished :arrow_down:0.0% :arrow_up:0.0%] [ec2-t3-xlarge-us-east-2](https://conbench.ursa.dev/compare/runs/3f8ac59fddef426cac31b6b2826ce5a6...e5b78c70b6fe4a699c4cb8b967fe5736/)
[Finished :arrow_down:0.38% :arrow_up:0.33%] [test-mac-arm](https://conbench.ursa.dev/compare/runs/02b74ab900d440f8b5361300f15ad15c...b870e78b617e41e98c2ade7241e0aa06/)
[Scheduled] [ursa-i9-9960x](https://conbench.ursa.dev/compare/runs/1f7740546747452aa61fece1295aa0e3...f66d8edd43724d9988628cbe4a4e4d50/)
[Finished :arrow_down:0.04% :arrow_up:0.26%] [ursa-thinkcentre-m75q](https://conbench.ursa.dev/compare/runs/c3c55d6999fe4d7eabec73cbdfa1fad6...4b9f45f6442447d3a9d758598d0e14d3/)
Buildkite builds:
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/415| `50fab73d` ec2-t3-xlarge-us-east-2>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-test-mac-arm/builds/401| `50fab73d` test-mac-arm>
[Scheduled] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-i9-9960x/builds/401| `50fab73d` ursa-i9-9960x>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-thinkcentre-m75q/builds/411| `50fab73d` ursa-thinkcentre-m75q>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/414| `0deefd18` ec2-t3-xlarge-us-east-2>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-test-mac-arm/builds/400| `0deefd18` test-mac-arm>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-i9-9960x/builds/400| `0deefd18` ursa-i9-9960x>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-thinkcentre-m75q/builds/410| `0deefd18` ursa-thinkcentre-m75q>
Supported benchmarks:
ec2-t3-xlarge-us-east-2: Supported benchmark langs: Python, R. Runs only benchmarks with cloud = True
test-mac-arm: Supported benchmark langs: C++, Python, R
ursa-i9-9960x: Supported benchmark langs: Python, R, JavaScript
ursa-thinkcentre-m75q: Supported benchmark langs: C++, Java
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829378845
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
Review comment:
i don't want all of this stuff to be heap allocated if it doesn't need to be (which it doesn't). i could do `std::array<const char *>` but that is also a pain because then i have to count the number of elements manually. i think they only added the template argument deduction for the size in C++17.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829549783
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
Review comment:
yes, and it matters a lot. these strings will be placed adjacent in memory (maybe spaced out for alignment) in the .rodata section of the binary. this means that if i load these strings for whatever reason (like generating comments!) it'll pull a bunch of them in within the same cache line, and while generating comments the whole block of strings will fit within cache.
what happens with `std::vector<std::string>`? before `main` is called, all of the constructors of static objects will be called. this means that we will allocate the vector on the heap. then, we'll do a million tiny allocations for these strings, and these strings will be copied from `.rodata` onto the heap. further, we have no idea where the allocator will put these strings, they could just be in random places in memory (i.e. terrible for cache).
so basically we end up with something that's worse in every way: double the memory gets used, cache performance is worse, startup time is longer
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829472744
##########
File path: cpp/src/arrow/compute/kernels/vector_sort.cc
##########
@@ -893,7 +893,8 @@ class TableSorter {
TableSorter(ExecContext* ctx, uint64_t* indices_begin, uint64_t* indices_end,
const Table& table, const SortOptions& options)
- : ctx_(ctx),
Review comment:
not in any very detailed way, just that it was inside this constructor, and there was some shared_ptr being deleted. it also didn't always reproduce, but it stopped reproducing once i made status be initialized before it got passed around to other stuff.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829355746
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
Review comment:
i switched it to use GetRandomSeed. in regards to nondeterminism between runs i'll let Weston decide.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1077327651
@westonpace Yes, I'm taking a look now. Thanks for reminding me.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r834009700
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row + num_offset;
+ int64_t num = 0;
+ int ibyte = static_cast<int>(num_offset);
+ for (; *num_str && ibyte < byte_width; ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ if (verify_padding) {
+ int num_chars = ibyte - num_offset;
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ VerifyStringAndNumber_Single(row, prefix, i, nums, byte_width, verify_padding);
+ }
+}
+
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ char tmp_str[256] = {};
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ std::memcpy(tmp_str, str + start, str_len);
+ VerifyStringAndNumber_Single(tmp_str, prefix, i, nums, sizeof(tmp_str),
+ verify_padding);
+ }
+}
+
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = off[i];
+ int32_t end = off[i + 1];
+ int32_t str_len = end - start;
+ ASSERT_LE(str_len, max_length);
+ ASSERT_GE(str_len, min_length);
+ for (int32_t i = start; i < end; i++) {
+ bool is_valid =
+ std::isdigit(str[i]) || std::isalpha(str[i]) || str[i] == ',' || str[i] == ' ';
+ ASSERT_TRUE(is_valid) << "Character " << str[i]
+ << " is not a digit, a letter, a comma, or a space";
+ }
+ }
+}
+
+void VerifyModuloBetween(const Datum& d, int32_t min, int32_t max, int32_t mod) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t m = n[i] % mod;
+ ASSERT_GE(m, min) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ ASSERT_LE(m, max) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ }
+}
+
+void VerifyAllBetween(const Datum& d, int32_t min, int32_t max) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ ASSERT_GE(n[i], min) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ ASSERT_LE(n[i], max) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ }
+}
+
+void VerifyNationKey(const Datum& d) { VerifyAllBetween(d, 0, 24); }
+
+void VerifyPhone(const Datum& d) {
+ int64_t length = d.length();
+ const char* phones = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ constexpr int kByteWidth = 15; // This is common for all PHONE columns
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = phones + i * kByteWidth;
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ }
+}
+
+void VerifyDecimalsBetween(const Datum& d, int64_t min, int64_t max) {
+ int64_t length = d.length();
+ const Decimal128* decs =
+ reinterpret_cast<const Decimal128*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int64_t val = static_cast<int64_t>(decs[i]);
+ ASSERT_LE(val, max);
+ ASSERT_GE(val, min);
+ }
+}
+
+void VerifyCorrectNumberOfWords_Varlen(const Datum& d, int num_words) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+
+ int32_t start = offsets[i];
+ int32_t end = offsets[i + 1];
+ int32_t str_len = end - start;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, str + start, str_len);
Review comment:
I see this resolved but you didn't make a change, is this deliberate?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082289658
I guess my biggest gripe is with the unnecessary heap allocation. But I guess it's not a big deal (until we begin measuring ExecPlans-per-second :stuck_out_tongue:)
I'll just revert the revert.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1057404174
And one more thing I though as I was doing this: For the tables that have keys shared across them (e.g. lineitems and orders), is it possible generate those at the same time so that the keys are the same across both? I can figure out if generating them separately is an issue by comparing duckdb or running the queries through arrow bench.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r817259916
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3738 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status Init();
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ void GenerateWord(size_t &offset, const char **words, size_t num_choices);
+ void GenerateNoun(size_t &offset);
+ void GenerateVerb(size_t &offset);
+ void GenerateAdjective(size_t &offset);
+ void GenerateAdverb(size_t &offset);
+ void GeneratePreposition(size_t &offset);
+ void GenerateAuxiliary(size_t &offset);
+ void GenerateTerminator(size_t &offset);
+
+ void GenerateNounPhrase(size_t &offset);
+ void GenerateVerbPhrase(size_t &offset);
+ void GeneratePrepositionalPhrase(size_t &offset);
+
+ void GenerateSentence(size_t &offset);
+
+ std::unique_ptr<Buffer> text_;
+ random::pcg32_fast rng_;
+ static constexpr size_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ int scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_generated_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
Review comment:
It's actually not really a copy - I have the `DCHECK`, made it work with 64-bit ints, return -1 in the unreachable scenario, and the variable and function names are different. I was mainly using it as a source for justifying why this is the fastest way.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane edited a comment on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane edited a comment on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1057480303
Though we _should_ be able to construct a plan that includes write nodes for both of those recordbatchreaders in C++ and then let C++ handle the rest... (no need to try and shoe horn that multithreading into R which will be reluctant to make that easy!) I'm digging into those now to see if I can wire that up myself...
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r824971804
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
Review comment:
Could you elaborate on this? I basically wanted to keep the "only abort once" logic inside of the Generator class, but the finished future belongs to the exec node.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r824984543
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string> &columns,
+ const std::vector<std::shared_ptr<DataType>> &types,
+ const std::unordered_map<std::string, int> &name_map,
+ std::vector<int> &gen_list)
+ {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if(columns.empty())
+ {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for(auto pair : name_map)
+ {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ }
+ else
+ {
+ for(const std::string &col : columns)
+ {
+ auto entry = name_map.find(col);
+ if(entry == name_map.end())
+ return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+ }
+
+ static TpchText g_text;
+
+ Status TpchText::InitIfNeeded(random::pcg32_fast &rng)
+ {
+ if(done_.load())
+ return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
Review comment:
Ah, I see now. All threads cooperatively allocate the 300MB array. That also explains my later question about why you are doing it by chunks. That allows you to divide up the work. I thought this was a "first caller is responsible for all initialization" scenario.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r817260152
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3738 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status Init();
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ void GenerateWord(size_t &offset, const char **words, size_t num_choices);
+ void GenerateNoun(size_t &offset);
+ void GenerateVerb(size_t &offset);
+ void GenerateAdjective(size_t &offset);
+ void GenerateAdverb(size_t &offset);
+ void GeneratePreposition(size_t &offset);
+ void GenerateAuxiliary(size_t &offset);
+ void GenerateTerminator(size_t &offset);
+
+ void GenerateNounPhrase(size_t &offset);
+ void GenerateVerbPhrase(size_t &offset);
+ void GeneratePrepositionalPhrase(size_t &offset);
+
+ void GenerateSentence(size_t &offset);
+
+ std::unique_ptr<Buffer> text_;
+ random::pcg32_fast rng_;
+ static constexpr size_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ int scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_generated_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
Review comment:
But yes, I'm not sure how different something like this has to be in order for it to count as "original" vs "copied and modified"...
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1057475963
Yeah, one thing to think about is what the R code would look like. You'd need some way of "linking" the two tables in R right?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1057569186
> Tangentially: ideally, I would like each table to go to a single parquet file (with multiple row groups, and written batch by batch). Do we have the node machinery for that already? Or just machinery for writing to datasets?
I'm not sure we need a dedicated write node for single-file writes.
What your describing is the "default partitioning" but looking at the code I'm not sure if we default to that or if we default to a segmentation fault when no partitioning is specified. :grimacing: `arrow::dataset::Partitioning::Default()` could be used in the meantime.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829345518
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row + num_offset;
+ int64_t num = 0;
+ int ibyte = static_cast<int>(num_offset);
+ for (; *num_str && ibyte < byte_width; ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ if (verify_padding) {
+ int num_chars = ibyte - num_offset;
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ VerifyStringAndNumber_Single(row, prefix, i, nums, byte_width, verify_padding);
+ }
+}
+
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ char tmp_str[256] = {};
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ std::memcpy(tmp_str, str + start, str_len);
+ VerifyStringAndNumber_Single(tmp_str, prefix, i, nums, sizeof(tmp_str),
+ verify_padding);
+ }
+}
+
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = off[i];
+ int32_t end = off[i + 1];
+ int32_t str_len = end - start;
+ ASSERT_LE(str_len, max_length);
+ ASSERT_GE(str_len, min_length);
+ for (int32_t i = start; i < end; i++) {
+ bool is_valid =
+ std::isdigit(str[i]) || std::isalpha(str[i]) || str[i] == ',' || str[i] == ' ';
+ ASSERT_TRUE(is_valid) << "Character " << str[i]
+ << " is not a digit, a letter, a comma, or a space";
+ }
+ }
+}
+
+void VerifyModuloBetween(const Datum& d, int32_t min, int32_t max, int32_t mod) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t m = n[i] % mod;
+ ASSERT_GE(m, min) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ ASSERT_LE(m, max) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ }
+}
+
+void VerifyAllBetween(const Datum& d, int32_t min, int32_t max) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ ASSERT_GE(n[i], min) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ ASSERT_LE(n[i], max) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ }
+}
+
+void VerifyNationKey(const Datum& d) { VerifyAllBetween(d, 0, 24); }
+
+void VerifyPhone(const Datum& d) {
+ int64_t length = d.length();
+ const char* phones = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ constexpr int kByteWidth = 15; // This is common for all PHONE columns
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = phones + i * kByteWidth;
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ }
+}
+
+void VerifyDecimalsBetween(const Datum& d, int64_t min, int64_t max) {
+ int64_t length = d.length();
+ const Decimal128* decs =
+ reinterpret_cast<const Decimal128*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int64_t val = static_cast<int64_t>(decs[i]);
+ ASSERT_LE(val, max);
+ ASSERT_GE(val, min);
+ }
+}
+
+void VerifyCorrectNumberOfWords_Varlen(const Datum& d, int num_words) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+
+ int32_t start = offsets[i];
+ int32_t end = offsets[i + 1];
+ int32_t str_len = end - start;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, str + start, str_len);
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = offsets[i]; j < offsets[i + 1]; j++) {
+ bool is_space = str[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(str[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << tmp_str;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << tmp_str;
+ }
+}
+
+void VerifyCorrectNumberOfWords_FixedWidth(const Datum& d, int num_words,
+ int byte_width) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+ const char* row = str + i * byte_width;
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = 0; j < byte_width && row[j]; j++) {
+ bool is_space = row[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(row[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << row;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << row;
+ }
+}
+
+void VerifyOneOf(const Datum& d, const std::unordered_set<char>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++)
+ ASSERT_TRUE(possibilities.find(col[i]) != possibilities.end());
+}
+
+void VerifyOneOf(const Datum& d, int32_t byte_width,
+ const std::unordered_set<std::string>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = col + i * byte_width;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, row, byte_width);
+ ASSERT_TRUE(possibilities.find(tmp_str) != possibilities.end())
+ << tmp_str << " is not a valid string.";
+ }
+}
+
+void CountInstances(std::unordered_map<int32_t, int32_t>& counts, const Datum& d) {
+ int64_t length = d.length();
+ const int32_t* nums = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) counts[nums[i]]++;
+}
+
+void CountModifiedComments(const Datum& d, int& good_count, int& bad_count) {
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ // Length of S_COMMENT is at most 100
+ char tmp_string[101];
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + offsets[i];
+ int32_t row_length = offsets[i + 1] - offsets[i];
+ std::memset(tmp_string, 0, sizeof(tmp_string));
+ std::memcpy(tmp_string, row, row_length);
+ char* customer = std::strstr(tmp_string, "Customer");
+ char* recommends = std::strstr(tmp_string, "Recommends");
+ char* complaints = std::strstr(tmp_string, "Complaints");
+ if (customer) {
+ ASSERT_TRUE((recommends != nullptr) ^ (complaints != nullptr));
+ if (recommends) good_count++;
+ if (complaints) bad_count++;
+ }
+ }
+}
+
+TEST(TpchNode, ScaleFactor) {
+ ExecContext ctx(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(&ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), 0.25f);
+ ExecNode* table = *gen.Supplier();
+ AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+ Declaration sink("sink", {Declaration::Input(table)}, SinkNodeOptions{&sink_gen});
+ std::ignore = *sink.AddToPlan(plan.get());
+ auto fut = StartAndCollect(plan.get(), sink_gen);
+ auto res = *fut.MoveResult();
+
+ int64_t kExpectedRows = 2500;
+ int64_t num_rows = 0;
+ for (auto& batch : res) num_rows += batch.length;
+ ASSERT_EQ(num_rows, kExpectedRows);
+ arrow::internal::GetCpuThreadPool()->WaitForIdle();
Review comment:
it was necessary at one point but i think it's not anymore. i've deleted it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829181919
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
Review comment:
Nit: use bracket around external dependencies.
```suggestion
#include <benchmark/benchmark.h>
```
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<float>(scale_factor));
+
+ ExecNode* lineitem =
+ *gen.Lineitem({"L_QUANTITY", "L_EXTENDEDPRICE", "L_TAX", "L_DISCOUNT", "L_SHIPDATE",
+ "L_RETURNFLAG", "L_LINESTATUS"});
+
+ std::shared_ptr<Date32Scalar> sept_2_1998 = std::make_shared<Date32Scalar>(
+ 10471); // September 2, 1998 is 10471 days after January 1, 1970
+ Expression filter =
+ less_equal(field_ref("L_SHIPDATE"), literal(std::move(sept_2_1998)));
+ FilterNodeOptions filter_opts(filter);
+
+ Expression l_returnflag = field_ref("L_RETURNFLAG");
+ Expression l_linestatus = field_ref("L_LINESTATUS");
+ Expression quantity = field_ref("L_QUANTITY");
+ Expression base_price = field_ref("L_EXTENDEDPRICE");
+
+ std::shared_ptr<Decimal128Scalar> decimal_1 =
+ std::make_shared<Decimal128Scalar>(Decimal128{0, 100}, decimal(12, 2));
+ Expression discount_multiplier =
+ call("subtract", {literal(decimal_1), field_ref("L_DISCOUNT")});
+ Expression tax_multiplier = call("add", {literal(decimal_1), field_ref("L_TAX")});
+ Expression disc_price =
+ call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier});
+ Expression charge =
+ call("multiply",
+ {call("cast",
+ {call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier})},
+ compute::CastOptions::Unsafe(decimal(12, 2))),
+ tax_multiplier});
+ Expression discount = field_ref("L_DISCOUNT");
+
+ std::vector<Expression> projection_list = {l_returnflag, l_linestatus, quantity,
+ base_price, disc_price, charge,
+ quantity, base_price, discount};
+ std::vector<std::string> project_names = {
+ "l_returnflag", "l_linestatus", "sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price", "avg_disc"};
+ ProjectNodeOptions project_opts(std::move(projection_list));
+
+ ScalarAggregateOptions sum_opts = ScalarAggregateOptions::Defaults();
+ CountOptions count_opts(CountOptions::CountMode::ALL);
+ std::vector<arrow::compute::internal::Aggregate> aggs = {
+ {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts},
+ {"hash_sum", &sum_opts}, {"hash_mean", &sum_opts}, {"hash_mean", &sum_opts},
+ {"hash_mean", &sum_opts}, {"hash_count", &count_opts}};
+
+ std::vector<FieldRef> cols = {2, 3, 4, 5, 6, 7, 8, 2};
+
+ std::vector<std::string> names = {"sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price",
+ "avg_disc", "count_order"};
+
+ std::vector<FieldRef> keys = {"L_RETURNFLAG", "L_LINESTATUS"};
+ AggregateNodeOptions agg_opts(aggs, cols, names, keys);
+
+ SortKey l_returnflag_key("L_RETURNFLAG");
+ SortKey l_linestatus_key("L_LINESTATUS");
+ SortOptions sort_opts({l_returnflag_key, l_linestatus_key});
+ OrderBySinkNodeOptions order_by_opts(sort_opts, sink_gen);
+
+ Declaration filter_decl("filter", {Declaration::Input(lineitem)}, filter_opts);
+ Declaration project_decl("project", project_opts);
+ Declaration aggregate_decl("aggregate", agg_opts);
+ Declaration orderby_decl("order_by_sink", order_by_opts);
+
+ Declaration q1 =
+ Declaration::Sequence({filter_decl, project_decl, aggregate_decl, orderby_decl});
+ std::ignore = *q1.AddToPlan(plan.get());
+ return plan;
+}
+
+static void BM_Tpch_Q1(benchmark::State& st) {
+ for (auto _ : st) {
+ st.PauseTiming();
+ AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+ std::shared_ptr<ExecPlan> plan = Plan_Q1(&sink_gen, static_cast<int>(st.range(0)));
+ st.ResumeTiming();
+ auto fut = StartAndCollect(plan.get(), sink_gen);
+ auto res = *fut.MoveResult();
+ }
+}
+
+BENCHMARK(BM_Tpch_Q1)->Args({1})->ArgNames({"SF"});
Review comment:
"SF" is for scale factor? Wouldn't it be nicer to spell it out in full?
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
+class OrdersAndLineItemGenerator;
+class PartAndPartSupplierGenerator;
+
+class ARROW_EXPORT TpchGen {
+ public:
+ /*
+ * \brief Create a factory for nodes that generate TPC-H data
+ *
+ * Note: Individual tables will reference each other. It is important that you only
+ * create a single TpchGen instance for each plan and then you can create nodes for each
+ * table from that single TpchGen instance. Note: Every batch will be scheduled as a new
+ * task using the ExecPlan's scheduler.
+ */
+ static Result<TpchGen> Make(ExecPlan* plan, float scale_factor = 1.0f,
+ int64_t batch_size = 4096);
Review comment:
Is it me, or does 4096 seem a bit small?
##########
File path: cpp/src/arrow/compute/kernels/vector_sort.cc
##########
@@ -893,7 +893,8 @@ class TableSorter {
TableSorter(ExecContext* ctx, uint64_t* indices_begin, uint64_t* indices_end,
const Table& table, const SortOptions& options)
- : ctx_(ctx),
+ : status_(),
+ ctx_(ctx),
Review comment:
Nit, but there's no need to explicitly spell out the default initialization for `status_`.
```suggestion
: ctx_(ctx),
```
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
Review comment:
Can you add a short comment at the beginning of each `Verify...` function explaining what it does?
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
Review comment:
Why aren't you using `util::string_view` or `std::string`?
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<float>(scale_factor));
+
+ ExecNode* lineitem =
+ *gen.Lineitem({"L_QUANTITY", "L_EXTENDEDPRICE", "L_TAX", "L_DISCOUNT", "L_SHIPDATE",
+ "L_RETURNFLAG", "L_LINESTATUS"});
+
+ std::shared_ptr<Date32Scalar> sept_2_1998 = std::make_shared<Date32Scalar>(
+ 10471); // September 2, 1998 is 10471 days after January 1, 1970
+ Expression filter =
+ less_equal(field_ref("L_SHIPDATE"), literal(std::move(sept_2_1998)));
+ FilterNodeOptions filter_opts(filter);
+
+ Expression l_returnflag = field_ref("L_RETURNFLAG");
+ Expression l_linestatus = field_ref("L_LINESTATUS");
+ Expression quantity = field_ref("L_QUANTITY");
+ Expression base_price = field_ref("L_EXTENDEDPRICE");
+
+ std::shared_ptr<Decimal128Scalar> decimal_1 =
+ std::make_shared<Decimal128Scalar>(Decimal128{0, 100}, decimal(12, 2));
+ Expression discount_multiplier =
+ call("subtract", {literal(decimal_1), field_ref("L_DISCOUNT")});
+ Expression tax_multiplier = call("add", {literal(decimal_1), field_ref("L_TAX")});
+ Expression disc_price =
+ call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier});
+ Expression charge =
+ call("multiply",
+ {call("cast",
+ {call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier})},
+ compute::CastOptions::Unsafe(decimal(12, 2))),
+ tax_multiplier});
+ Expression discount = field_ref("L_DISCOUNT");
+
+ std::vector<Expression> projection_list = {l_returnflag, l_linestatus, quantity,
+ base_price, disc_price, charge,
+ quantity, base_price, discount};
+ std::vector<std::string> project_names = {
+ "l_returnflag", "l_linestatus", "sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price", "avg_disc"};
+ ProjectNodeOptions project_opts(std::move(projection_list));
+
+ ScalarAggregateOptions sum_opts = ScalarAggregateOptions::Defaults();
+ CountOptions count_opts(CountOptions::CountMode::ALL);
+ std::vector<arrow::compute::internal::Aggregate> aggs = {
+ {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts},
+ {"hash_sum", &sum_opts}, {"hash_mean", &sum_opts}, {"hash_mean", &sum_opts},
+ {"hash_mean", &sum_opts}, {"hash_count", &count_opts}};
+
+ std::vector<FieldRef> cols = {2, 3, 4, 5, 6, 7, 8, 2};
Review comment:
Without changing anything functionally, perhaps it would be nicer to initialize this vector from field names rather than field indices?
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<float>(scale_factor));
Review comment:
Why `float`? This seems unimportant, but should probably pass a `double` for future-proofing.
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row + num_offset;
+ int64_t num = 0;
+ int ibyte = static_cast<int>(num_offset);
+ for (; *num_str && ibyte < byte_width; ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ if (verify_padding) {
+ int num_chars = ibyte - num_offset;
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ VerifyStringAndNumber_Single(row, prefix, i, nums, byte_width, verify_padding);
+ }
+}
+
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ char tmp_str[256] = {};
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ std::memcpy(tmp_str, str + start, str_len);
+ VerifyStringAndNumber_Single(tmp_str, prefix, i, nums, sizeof(tmp_str),
+ verify_padding);
+ }
+}
+
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = off[i];
+ int32_t end = off[i + 1];
+ int32_t str_len = end - start;
+ ASSERT_LE(str_len, max_length);
+ ASSERT_GE(str_len, min_length);
+ for (int32_t i = start; i < end; i++) {
+ bool is_valid =
+ std::isdigit(str[i]) || std::isalpha(str[i]) || str[i] == ',' || str[i] == ' ';
+ ASSERT_TRUE(is_valid) << "Character " << str[i]
+ << " is not a digit, a letter, a comma, or a space";
+ }
+ }
+}
+
+void VerifyModuloBetween(const Datum& d, int32_t min, int32_t max, int32_t mod) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t m = n[i] % mod;
+ ASSERT_GE(m, min) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ ASSERT_LE(m, max) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ }
+}
+
+void VerifyAllBetween(const Datum& d, int32_t min, int32_t max) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ ASSERT_GE(n[i], min) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ ASSERT_LE(n[i], max) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ }
+}
+
+void VerifyNationKey(const Datum& d) { VerifyAllBetween(d, 0, 24); }
+
+void VerifyPhone(const Datum& d) {
+ int64_t length = d.length();
+ const char* phones = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ constexpr int kByteWidth = 15; // This is common for all PHONE columns
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = phones + i * kByteWidth;
Review comment:
It's probably allowed to use `std::regex` in tests (I see we already use it in some dataset tests).
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
Review comment:
Why is `pcg_random.h` included here? Generally we should avoid including stuff in `.h` files if it's only needed for the implementations in the corresponding `.cc`.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
Review comment:
Why doesn't this use the existing formatting functionalities in `arrow/util/formatting.h`?
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ if (!partsupp_output_queue_.empty()) {
+ ExecBatch result = std::move(partsupp_output_queue_.front());
+ partsupp_output_queue_.pop();
+ return std::move(result);
+ }
+ }
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+ if (!part_cols_.empty()) {
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch part_batch,
+ ExecBatch::Make(std::move(part_result)));
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ part_output_queue_.emplace(std::move(part_batch));
+ }
+ }
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ for (size_t i = 1; i < partsupp_results.size(); i++)
+ partsupp_output_queue_.emplace(std::move(partsupp_results[i]));
+ }
+ return std::move(partsupp_results[0]);
+ }
+
+ private:
+#define FOR_EACH_PART_COLUMN(F) \
+ F(P_PARTKEY) \
+ F(P_NAME) \
+ F(P_MFGR) \
+ F(P_BRAND) \
+ F(P_TYPE) \
+ F(P_SIZE) \
+ F(P_CONTAINER) \
+ F(P_RETAILPRICE) \
+ F(P_COMMENT)
+
+#define FOR_EACH_PARTSUPP_COLUMN(F) \
+ F(PS_PARTKEY) \
+ F(PS_SUPPKEY) \
+ F(PS_AVAILQTY) \
+ F(PS_SUPPLYCOST) \
+ F(PS_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct PART {
+ enum {
+ FOR_EACH_PART_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct PARTSUPP {
+ enum {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, PART::col},
+ const std::unordered_map<std::string, int> kPartNameMap = {
+ FOR_EACH_PART_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, PARTSUPP::col},
+ const std::unordered_map<std::string, int> kPartsuppNameMap = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kPartGenerators = {FOR_EACH_PART_COLUMN(MAKE_FN_ARRAY)};
+ std::vector<GenerateColumnFn> kPartsuppGenerators = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kPartTypes = {
+ int32(), utf8(), fixed_size_binary(25), fixed_size_binary(10),
+ utf8(), int32(), fixed_size_binary(10), decimal(12, 2),
+ utf8(),
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kPartsuppTypes = {
+ int32(), int32(), int32(), decimal(12, 2), utf8(),
+ };
+
+ Status AllocatePartBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.part[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.part_to_generate * byte_width));
+ ArrayData ad(kPartTypes[column], tld.part_to_generate, {nullptr, std::move(buff)});
+ tld.part[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status P_PARTKEY(size_t thread_index) {
Review comment:
Can you add comments describing and explaining the various generation methods below?
We should be able to understand and maintain this code without knowing the TPC-H spec by heart, ideally :-)
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
+class OrdersAndLineItemGenerator;
+class PartAndPartSupplierGenerator;
+
+class ARROW_EXPORT TpchGen {
+ public:
+ /*
+ * \brief Create a factory for nodes that generate TPC-H data
+ *
+ * Note: Individual tables will reference each other. It is important that you only
+ * create a single TpchGen instance for each plan and then you can create nodes for each
+ * table from that single TpchGen instance. Note: Every batch will be scheduled as a new
+ * task using the ExecPlan's scheduler.
+ */
+ static Result<TpchGen> Make(ExecPlan* plan, float scale_factor = 1.0f,
+ int64_t batch_size = 4096);
+
+ Result<ExecNode*> Supplier(std::vector<std::string> columns = {});
Review comment:
Can you add a terse description for these methods, so that non-TPC-H experts don't have to guess too much? (at least a one-liner)
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<float>(scale_factor));
+
+ ExecNode* lineitem =
+ *gen.Lineitem({"L_QUANTITY", "L_EXTENDEDPRICE", "L_TAX", "L_DISCOUNT", "L_SHIPDATE",
+ "L_RETURNFLAG", "L_LINESTATUS"});
+
+ std::shared_ptr<Date32Scalar> sept_2_1998 = std::make_shared<Date32Scalar>(
+ 10471); // September 2, 1998 is 10471 days after January 1, 1970
Review comment:
Nit, but can use `auto` to avoid repeating a long type and `const` to stress that this won't be mutated:
```suggestion
const auto sept_2_1998 = std::make_shared<Date32Scalar>(
10471); // September 2, 1998 is 10471 days after January 1, 1970
```
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
+class OrdersAndLineItemGenerator;
+class PartAndPartSupplierGenerator;
+
+class ARROW_EXPORT TpchGen {
+ public:
+ /*
+ * \brief Create a factory for nodes that generate TPC-H data
+ *
+ * Note: Individual tables will reference each other. It is important that you only
+ * create a single TpchGen instance for each plan and then you can create nodes for each
+ * table from that single TpchGen instance. Note: Every batch will be scheduled as a new
+ * task using the ExecPlan's scheduler.
+ */
+ static Result<TpchGen> Make(ExecPlan* plan, float scale_factor = 1.0f,
Review comment:
Nit: there doesn't seem to be a reason to pass a `float` here, just use a `double`?
##########
File path: cpp/src/arrow/vendored/CMakeLists.txt
##########
@@ -19,3 +19,4 @@ arrow_install_all_headers("arrow/vendored")
add_subdirectory(datetime)
add_subdirectory(double-conversion)
+add_subdirectory(pcg)
Review comment:
I _think_ we should be able to revert these changes if you're not including the PCG headers from a `.h`.
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
Review comment:
Please try to avoid including `api.h` files as they're generally costly.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
Review comment:
Put all this in the `arrow::compute::internal` namespace?
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
Review comment:
Mutable argument should be passed by pointer, not reference. Also usually they will be the last argument by position:
```suggestion
void VerifyUniqueKey(const Datum& d, int32_t min, int32_t max,
std::unordered_set<int32_t>* seen) {
```
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
Review comment:
Hmm, generally you should call the `Validate` or `ValidateFull` method on the object, rather than directly call functions from `validate.h`.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
Review comment:
Also, I'm curious: is this deliberately non-deterministic between runs? @westonpace
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row + num_offset;
+ int64_t num = 0;
+ int ibyte = static_cast<int>(num_offset);
+ for (; *num_str && ibyte < byte_width; ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
Review comment:
What is this doing? Please at least add a comment explaining. But it also seems this might be made higher level instead of parsing by hand?
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row + num_offset;
+ int64_t num = 0;
+ int ibyte = static_cast<int>(num_offset);
+ for (; *num_str && ibyte < byte_width; ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ if (verify_padding) {
+ int num_chars = ibyte - num_offset;
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ VerifyStringAndNumber_Single(row, prefix, i, nums, byte_width, verify_padding);
+ }
+}
+
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ char tmp_str[256] = {};
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ std::memcpy(tmp_str, str + start, str_len);
+ VerifyStringAndNumber_Single(tmp_str, prefix, i, nums, sizeof(tmp_str),
+ verify_padding);
+ }
+}
+
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
Review comment:
Same comment here and for all other Verify functions about visiting using higher-level constructs such as `VisitArrayDataInline`.
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row + num_offset;
+ int64_t num = 0;
+ int ibyte = static_cast<int>(num_offset);
+ for (; *num_str && ibyte < byte_width; ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ if (verify_padding) {
+ int num_chars = ibyte - num_offset;
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ VerifyStringAndNumber_Single(row, prefix, i, nums, byte_width, verify_padding);
+ }
+}
+
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ char tmp_str[256] = {};
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ std::memcpy(tmp_str, str + start, str_len);
+ VerifyStringAndNumber_Single(tmp_str, prefix, i, nums, sizeof(tmp_str),
+ verify_padding);
+ }
Review comment:
You can visit a strings array using e.g. `VisitArrayDataInline`. No need to inspect all buffers by hand. Also it seems you're not checking for null values?
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row + num_offset;
+ int64_t num = 0;
+ int ibyte = static_cast<int>(num_offset);
+ for (; *num_str && ibyte < byte_width; ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ if (verify_padding) {
+ int num_chars = ibyte - num_offset;
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ VerifyStringAndNumber_Single(row, prefix, i, nums, byte_width, verify_padding);
+ }
+}
+
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ char tmp_str[256] = {};
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ std::memcpy(tmp_str, str + start, str_len);
+ VerifyStringAndNumber_Single(tmp_str, prefix, i, nums, sizeof(tmp_str),
+ verify_padding);
+ }
+}
+
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = off[i];
+ int32_t end = off[i + 1];
+ int32_t str_len = end - start;
+ ASSERT_LE(str_len, max_length);
+ ASSERT_GE(str_len, min_length);
+ for (int32_t i = start; i < end; i++) {
+ bool is_valid =
+ std::isdigit(str[i]) || std::isalpha(str[i]) || str[i] == ',' || str[i] == ' ';
+ ASSERT_TRUE(is_valid) << "Character " << str[i]
+ << " is not a digit, a letter, a comma, or a space";
+ }
+ }
+}
+
+void VerifyModuloBetween(const Datum& d, int32_t min, int32_t max, int32_t mod) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t m = n[i] % mod;
+ ASSERT_GE(m, min) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ ASSERT_LE(m, max) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ }
+}
+
+void VerifyAllBetween(const Datum& d, int32_t min, int32_t max) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ ASSERT_GE(n[i], min) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ ASSERT_LE(n[i], max) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ }
+}
+
+void VerifyNationKey(const Datum& d) { VerifyAllBetween(d, 0, 24); }
+
+void VerifyPhone(const Datum& d) {
+ int64_t length = d.length();
+ const char* phones = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ constexpr int kByteWidth = 15; // This is common for all PHONE columns
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = phones + i * kByteWidth;
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ }
+}
+
+void VerifyDecimalsBetween(const Datum& d, int64_t min, int64_t max) {
+ int64_t length = d.length();
+ const Decimal128* decs =
+ reinterpret_cast<const Decimal128*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int64_t val = static_cast<int64_t>(decs[i]);
+ ASSERT_LE(val, max);
+ ASSERT_GE(val, min);
+ }
+}
+
+void VerifyCorrectNumberOfWords_Varlen(const Datum& d, int num_words) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+
+ int32_t start = offsets[i];
+ int32_t end = offsets[i + 1];
+ int32_t str_len = end - start;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, str + start, str_len);
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = offsets[i]; j < offsets[i + 1]; j++) {
+ bool is_space = str[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(str[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << tmp_str;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << tmp_str;
+ }
+}
+
+void VerifyCorrectNumberOfWords_FixedWidth(const Datum& d, int num_words,
+ int byte_width) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+ const char* row = str + i * byte_width;
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = 0; j < byte_width && row[j]; j++) {
+ bool is_space = row[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(row[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << row;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << row;
+ }
+}
+
+void VerifyOneOf(const Datum& d, const std::unordered_set<char>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++)
+ ASSERT_TRUE(possibilities.find(col[i]) != possibilities.end());
+}
+
+void VerifyOneOf(const Datum& d, int32_t byte_width,
+ const std::unordered_set<std::string>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = col + i * byte_width;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, row, byte_width);
+ ASSERT_TRUE(possibilities.find(tmp_str) != possibilities.end())
+ << tmp_str << " is not a valid string.";
+ }
+}
+
+void CountInstances(std::unordered_map<int32_t, int32_t>& counts, const Datum& d) {
+ int64_t length = d.length();
+ const int32_t* nums = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) counts[nums[i]]++;
+}
+
+void CountModifiedComments(const Datum& d, int& good_count, int& bad_count) {
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ // Length of S_COMMENT is at most 100
+ char tmp_string[101];
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + offsets[i];
+ int32_t row_length = offsets[i + 1] - offsets[i];
+ std::memset(tmp_string, 0, sizeof(tmp_string));
+ std::memcpy(tmp_string, row, row_length);
+ char* customer = std::strstr(tmp_string, "Customer");
+ char* recommends = std::strstr(tmp_string, "Recommends");
+ char* complaints = std::strstr(tmp_string, "Complaints");
+ if (customer) {
+ ASSERT_TRUE((recommends != nullptr) ^ (complaints != nullptr));
+ if (recommends) good_count++;
+ if (complaints) bad_count++;
+ }
+ }
+}
+
+TEST(TpchNode, ScaleFactor) {
+ ExecContext ctx(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(&ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), 0.25f);
+ ExecNode* table = *gen.Supplier();
+ AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+ Declaration sink("sink", {Declaration::Input(table)}, SinkNodeOptions{&sink_gen});
+ std::ignore = *sink.AddToPlan(plan.get());
+ auto fut = StartAndCollect(plan.get(), sink_gen);
+ auto res = *fut.MoveResult();
+
+ int64_t kExpectedRows = 2500;
+ int64_t num_rows = 0;
+ for (auto& batch : res) num_rows += batch.length;
+ ASSERT_EQ(num_rows, kExpectedRows);
+ arrow::internal::GetCpuThreadPool()->WaitForIdle();
Review comment:
Why is this necessary? Can you add a comment here?
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row + num_offset;
+ int64_t num = 0;
+ int ibyte = static_cast<int>(num_offset);
+ for (; *num_str && ibyte < byte_width; ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ if (verify_padding) {
+ int num_chars = ibyte - num_offset;
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ VerifyStringAndNumber_Single(row, prefix, i, nums, byte_width, verify_padding);
+ }
+}
+
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ char tmp_str[256] = {};
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ std::memcpy(tmp_str, str + start, str_len);
+ VerifyStringAndNumber_Single(tmp_str, prefix, i, nums, sizeof(tmp_str),
+ verify_padding);
+ }
+}
+
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = off[i];
+ int32_t end = off[i + 1];
+ int32_t str_len = end - start;
+ ASSERT_LE(str_len, max_length);
+ ASSERT_GE(str_len, min_length);
+ for (int32_t i = start; i < end; i++) {
+ bool is_valid =
+ std::isdigit(str[i]) || std::isalpha(str[i]) || str[i] == ',' || str[i] == ' ';
+ ASSERT_TRUE(is_valid) << "Character " << str[i]
+ << " is not a digit, a letter, a comma, or a space";
+ }
+ }
+}
+
+void VerifyModuloBetween(const Datum& d, int32_t min, int32_t max, int32_t mod) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t m = n[i] % mod;
+ ASSERT_GE(m, min) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ ASSERT_LE(m, max) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ }
+}
+
+void VerifyAllBetween(const Datum& d, int32_t min, int32_t max) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ ASSERT_GE(n[i], min) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ ASSERT_LE(n[i], max) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ }
+}
+
+void VerifyNationKey(const Datum& d) { VerifyAllBetween(d, 0, 24); }
+
+void VerifyPhone(const Datum& d) {
+ int64_t length = d.length();
+ const char* phones = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ constexpr int kByteWidth = 15; // This is common for all PHONE columns
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = phones + i * kByteWidth;
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ }
+}
+
+void VerifyDecimalsBetween(const Datum& d, int64_t min, int64_t max) {
+ int64_t length = d.length();
+ const Decimal128* decs =
+ reinterpret_cast<const Decimal128*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int64_t val = static_cast<int64_t>(decs[i]);
+ ASSERT_LE(val, max);
+ ASSERT_GE(val, min);
+ }
+}
+
+void VerifyCorrectNumberOfWords_Varlen(const Datum& d, int num_words) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+
+ int32_t start = offsets[i];
+ int32_t end = offsets[i + 1];
+ int32_t str_len = end - start;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, str + start, str_len);
Review comment:
Please let's refrain from using C-style code, you can use e.g. `util::string_view` to handle a range of bytes as a string.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
Review comment:
Hmm, unless there's a strong reason to have a module global like this, it is generally better to have a local variable or member attribute where needed.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
Review comment:
Can you put all private declarations in the anonymous space? That will hide them from outside code and then you don't have to add `static` to some declarations.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ if (!partsupp_output_queue_.empty()) {
+ ExecBatch result = std::move(partsupp_output_queue_.front());
+ partsupp_output_queue_.pop();
+ return std::move(result);
+ }
+ }
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+ if (!part_cols_.empty()) {
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch part_batch,
+ ExecBatch::Make(std::move(part_result)));
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ part_output_queue_.emplace(std::move(part_batch));
+ }
+ }
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ for (size_t i = 1; i < partsupp_results.size(); i++)
+ partsupp_output_queue_.emplace(std::move(partsupp_results[i]));
+ }
+ return std::move(partsupp_results[0]);
+ }
+
+ private:
+#define FOR_EACH_PART_COLUMN(F) \
+ F(P_PARTKEY) \
+ F(P_NAME) \
+ F(P_MFGR) \
+ F(P_BRAND) \
+ F(P_TYPE) \
+ F(P_SIZE) \
+ F(P_CONTAINER) \
+ F(P_RETAILPRICE) \
+ F(P_COMMENT)
+
+#define FOR_EACH_PARTSUPP_COLUMN(F) \
+ F(PS_PARTKEY) \
+ F(PS_SUPPKEY) \
+ F(PS_AVAILQTY) \
+ F(PS_SUPPLYCOST) \
+ F(PS_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct PART {
+ enum {
+ FOR_EACH_PART_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct PARTSUPP {
+ enum {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, PART::col},
+ const std::unordered_map<std::string, int> kPartNameMap = {
+ FOR_EACH_PART_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, PARTSUPP::col},
+ const std::unordered_map<std::string, int> kPartsuppNameMap = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kPartGenerators = {FOR_EACH_PART_COLUMN(MAKE_FN_ARRAY)};
+ std::vector<GenerateColumnFn> kPartsuppGenerators = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kPartTypes = {
+ int32(), utf8(), fixed_size_binary(25), fixed_size_binary(10),
+ utf8(), int32(), fixed_size_binary(10), decimal(12, 2),
+ utf8(),
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kPartsuppTypes = {
+ int32(), int32(), int32(), decimal(12, 2), utf8(),
+ };
+
+ Status AllocatePartBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.part[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.part_to_generate * byte_width));
+ ArrayData ad(kPartTypes[column], tld.part_to_generate, {nullptr, std::move(buff)});
+ tld.part[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status P_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_PARTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_PARTKEY));
+ int32_t* p_partkey = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) {
+ p_partkey[i] = static_cast<int32_t>(tld.partkey_start + i + 1);
+ ARROW_DCHECK(1 <= p_partkey[i] && p_partkey[i] <= part_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_NAME].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(
+ 0, static_cast<uint8_t>(kNumNameParts - 1));
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dist(tld.rng));
+ tld.string_indices[irow * 5 + ipart] = name_part_index;
+ string_length += std::strlen(NameParts[name_part_index]);
+ }
+ // Add 4 because there is a space between each word (i.e. four spaces)
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length + 4);
+ }
+ // Add an extra byte for the space after in the very last string.
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate] + 1));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 5 + ipart];
+ const char* part = NameParts[name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ *row++ = ' ';
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_NAME], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_NAME] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_MFGR(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_MFGR].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* manufacturer = "Manufacturer#";
+ const size_t manufacturer_length = std::strlen(manufacturer);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_MFGR));
+ char* p_mfgr = reinterpret_cast<char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ std::strncpy(p_mfgr + byte_width * irow, manufacturer, byte_width);
+ char mfgr_id = '0' + dist(tld.rng);
+ *(p_mfgr + byte_width * irow + manufacturer_length) = mfgr_id;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_BRAND(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_BRAND].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_MFGR(thread_index));
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* brand = "Brand#";
+ const size_t brand_length = std::strlen(brand);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_BRAND));
+ const char* p_mfgr = reinterpret_cast<const char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->data());
+ char* p_brand = reinterpret_cast<char*>(
+ tld.part[PART::P_BRAND].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_BRAND]);
+ int32_t mfgr_byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ const size_t mfgr_id_offset = std::strlen("Manufacturer#");
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = p_brand + byte_width * irow;
+ char mfgr_id = *(p_mfgr + irow * mfgr_byte_width + mfgr_id_offset);
+ char brand_id = '0' + dist(tld.rng);
+ std::strncpy(row, brand, byte_width);
+ *(row + brand_length) = mfgr_id;
+ *(row + brand_length + 1) = brand_id;
+ irow += 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_TYPE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_TYPE].kind() == Datum::NONE) {
+ using D = std::uniform_int_distribution<uint32_t>;
+ D dists[] = {
+ D{0, static_cast<uint8_t>(kNumTypes_1 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_2 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_3 - 1)},
+ };
+
+ const char** types[] = {Types_1, Types_2, Types_3};
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dists[ipart](tld.rng));
+ tld.string_indices[irow * 3 + ipart] = name_part_index;
+ string_length += std::strlen(types[ipart][name_part_index]);
+ }
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate]));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 3 + ipart];
+ const char* part = types[ipart][name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_TYPE], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_TYPE] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_SIZE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_SIZE].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(1, 50);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_SIZE));
+ int32_t* p_size = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_SIZE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) p_size[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status P_CONTAINER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_CONTAINER].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist1(
+ 0, static_cast<uint8_t>(kNumContainers_1 - 1));
+ std::uniform_int_distribution<int> dist2(
+ 0, static_cast<uint8_t>(kNumContainers_2 - 1));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_CONTAINER));
+ char* p_container = reinterpret_cast<char*>(
+ tld.part[PART::P_CONTAINER].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_CONTAINER]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int container1_idx = dist1(tld.rng);
+ int container2_idx = dist2(tld.rng);
+ const char* container1 = Containers_1[container1_idx];
+ const char* container2 = Containers_2[container2_idx];
+ size_t container1_length = std::strlen(container1);
+ size_t container2_length = std::strlen(container2);
+
+ char* row = p_container + byte_width * irow;
+ std::strncpy(row, container1, byte_width);
+ std::memcpy(row + container1_length, container2, container2_length);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_RETAILPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_RETAILPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_RETAILPRICE));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+ Decimal128* p_retailprice = reinterpret_cast<Decimal128*>(
+ tld.part[PART::P_RETAILPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int32_t partkey = p_partkey[irow];
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ p_retailprice[irow] = {retail_price};
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.part[PART::P_COMMENT],
+ g_text.GenerateComments(tld.part_to_generate, 5, 22, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ int64_t PartsuppBatchesToGenerate(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t num_batches = (ps_to_generate + batch_size_ - 1) / batch_size_;
+ return num_batches;
+ }
+
+ Status InitPartsupp(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.generated_partsupp.reset();
+ int64_t num_batches = PartsuppBatchesToGenerate(thread_index);
+ tld.partsupp.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.partsupp) {
+ batch.resize(PARTSUPP::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocatePartSuppBatch(size_t thread_index, size_t ibatch, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartsuppTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kPartsuppTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.partsupp[ibatch][column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status PS_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_PARTKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_PARTKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_PARTKEY));
+ int32_t* ps_partkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run; ipartsupp++, irun++)
+ ps_partkey[batch_offset++] = p_partkey[ipart];
+
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPKEY));
+ int32_t* ps_suppkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run;
+ ipartsupp++, irun++) {
+ int32_t supplier = static_cast<int32_t>(ipartsupp);
+ int32_t partkey = p_partkey[ipart];
+ ps_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_AVAILQTY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_AVAILQTY]) {
+ tld.generated_partsupp[PARTSUPP::PS_AVAILQTY] = true;
+ std::uniform_int_distribution<int32_t> dist(1, 9999);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_AVAILQTY));
+ int32_t* ps_availqty =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for (int64_t irun = 0; irun < next_run; irun++) ps_availqty[irun] = dist(tld.rng);
+
+ tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPLYCOST(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST] = true;
+ std::uniform_int_distribution<int64_t> dist(100, 100000);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(
+ AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPLYCOST));
+ Decimal128* ps_supplycost =
+ reinterpret_cast<Decimal128*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for (int64_t irun = 0; irun < next_run; irun++)
+ ps_supplycost[irun] = {dist(tld.rng)};
+
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PARTSUPP::PS_COMMENT].kind() == Datum::NONE) {
+ int64_t irow = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ int64_t num_rows = std::min(batch_size_, ps_to_generate - irow);
+ ARROW_ASSIGN_OR_RAISE(tld.partsupp[ibatch][PARTSUPP::PS_COMMENT],
+ g_text.GenerateComments(num_rows, 49, 198, tld.rng));
+ irow += num_rows;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ std::vector<Datum> part;
+ std::vector<int8_t> string_indices;
+ int64_t part_to_generate{0};
+ int64_t partkey_start{0};
+
+ std::vector<std::vector<Datum>> partsupp;
+ std::bitset<PARTSUPP::kNumCols> generated_partsupp;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex part_output_queue_mutex_;
+ std::mutex partsupp_output_queue_mutex_;
+ std::queue<ExecBatch> part_output_queue_;
+ std::queue<ExecBatch> partsupp_output_queue_;
+ int64_t batch_size_{0};
+ float scale_factor_{0};
+ int64_t part_rows_to_generate_{0};
+ int64_t part_rows_generated_{0};
+ std::vector<int> part_cols_;
+ std::vector<int> partsupp_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<int64_t> part_batches_generated_ = {0};
+ std::atomic<int64_t> partsupp_batches_generated_ = {0};
+ static constexpr int64_t kPartSuppRowsPerPart = 4;
+};
+
+class OrdersAndLineItemGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ tld.items_per_order.resize(batch_size_);
+ tld.rng.seed(seq);
+ }
+ orders_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000 * 10);
+ }
+ return Status::OK();
+ }
+
+ int64_t orders_batches_generated() const { return orders_batches_generated_.load(); }
+
+ int64_t lineitem_batches_generated() const {
+ return lineitem_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetOrdersOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kOrdersTypes, kOrdersNameMap, orders_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetLineItemOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kLineitemTypes, kLineitemNameMap, lineitem_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextOrdersBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if (!orders_output_queue_.empty()) {
+ ExecBatch batch = std::move(orders_output_queue_.front());
+ orders_output_queue_.pop();
+ return std::move(batch);
+ } else if (orders_rows_generated_ == orders_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate =
+ std::min(batch_size_, orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.first_batch_offset = 0;
+ tld.generated_lineitem.reset();
+
+ for (int col : orders_cols_) RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for (int col : lineitem_cols_) RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for (size_t i = 0; i < orders_cols_.size(); i++) {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ if (!lineitem_cols_.empty()) {
+ std::vector<ExecBatch> lineitem_results;
+ for (size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++) {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for (size_t icol = 0; icol < lineitem_cols_.size(); icol++) {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(lineitem_output_queue_mutex_);
+ for (ExecBatch& eb : lineitem_results) {
+ lineitem_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(orders_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextLineItemBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ExecBatch queued;
+ bool from_queue = false;
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ if (!lineitem_output_queue_.empty()) {
+ queued = std::move(lineitem_output_queue_.front());
+ lineitem_output_queue_.pop();
+ from_queue = true;
+ }
+ }
+ tld.first_batch_offset = 0;
+ if (from_queue) {
+ ARROW_DCHECK(queued.length <= batch_size_);
+ tld.first_batch_offset = queued.length;
+ if (queued.length == batch_size_) return std::move(queued);
+ }
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if (orders_rows_generated_ == orders_rows_to_generate_) {
+ if (from_queue) return std::move(queued);
+ return util::nullopt;
+ }
+
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate =
+ std::min(batch_size_, orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1ll);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.generated_lineitem.reset();
+ if (from_queue) {
+ lineitem_batches_generated_.fetch_sub(1);
+ for (size_t i = 0; i < lineitem_cols_.size(); i++)
+ if (tld.lineitem[0][lineitem_cols_[i]].kind() == Datum::NONE)
+ tld.lineitem[0][lineitem_cols_[i]] = std::move(queued[i]);
+ }
+
+ for (int col : orders_cols_) RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for (int col : lineitem_cols_) RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ if (!orders_cols_.empty()) {
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for (size_t i = 0; i < orders_cols_.size(); i++) {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch orders_batch,
+ ExecBatch::Make(std::move(orders_result)));
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ orders_output_queue_.emplace(std::move(orders_batch));
+ }
+ }
+ std::vector<ExecBatch> lineitem_results;
+ for (size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++) {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for (size_t icol = 0; icol < lineitem_cols_.size(); icol++) {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ lineitem_batches_generated_.fetch_add(static_cast<int64_t>(lineitem_results.size()));
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ for (size_t i = 1; i < lineitem_results.size(); i++)
+ lineitem_output_queue_.emplace(std::move(lineitem_results[i]));
+ }
+ return std::move(lineitem_results[0]);
+ }
+
+ private:
+#define FOR_EACH_ORDERS_COLUMN(F) \
+ F(O_ORDERKEY) \
+ F(O_CUSTKEY) \
+ F(O_ORDERSTATUS) \
+ F(O_TOTALPRICE) \
+ F(O_ORDERDATE) \
+ F(O_ORDERPRIORITY) \
+ F(O_CLERK) \
+ F(O_SHIPPRIORITY) \
+ F(O_COMMENT)
+
+#define FOR_EACH_LINEITEM_COLUMN(F) \
+ F(L_ORDERKEY) \
+ F(L_PARTKEY) \
+ F(L_SUPPKEY) \
+ F(L_LINENUMBER) \
+ F(L_QUANTITY) \
+ F(L_EXTENDEDPRICE) \
+ F(L_DISCOUNT) \
+ F(L_TAX) \
+ F(L_RETURNFLAG) \
+ F(L_LINESTATUS) \
+ F(L_SHIPDATE) \
+ F(L_COMMITDATE) \
+ F(L_RECEIPTDATE) \
+ F(L_SHIPINSTRUCT) \
+ F(L_SHIPMODE) \
+ F(L_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct ORDERS {
+ enum {
+ FOR_EACH_ORDERS_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct LINEITEM {
+ enum {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, ORDERS::col},
+ const std::unordered_map<std::string, int> kOrdersNameMap = {
+ FOR_EACH_ORDERS_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, LINEITEM::col},
+ const std::unordered_map<std::string, int> kLineitemNameMap = {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ const std::vector<GenerateColumnFn> kOrdersGenerators = {
+ FOR_EACH_ORDERS_COLUMN(MAKE_FN_ARRAY)};
+ const std::vector<GenerateColumnFn> kLineitemGenerators = {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kOrdersTypes = {int32(),
+ int32(),
+ fixed_size_binary(1),
+ decimal(12, 2),
+ date32(),
+ fixed_size_binary(15),
+ fixed_size_binary(15),
+ int32(),
+ utf8()};
+
+ const std::vector<std::shared_ptr<DataType>> kLineitemTypes = {
+ int32(),
+ int32(),
+ int32(),
+ int32(),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ fixed_size_binary(1),
+ fixed_size_binary(1),
+ date32(),
+ date32(),
+ date32(),
+ fixed_size_binary(25),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status AllocateOrdersBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.orders[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.orders_to_generate * byte_width));
+ ArrayData ad(kOrdersTypes[column], tld.orders_to_generate,
+ {nullptr, std::move(buff)});
+ tld.orders[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status O_ORDERKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERKEY));
+ int32_t* o_orderkey = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ int32_t orderkey_index = static_cast<int32_t>(tld.orderkey_start + i);
+ int32_t index_of_run = orderkey_index / 8;
+ int32_t index_in_run = orderkey_index % 8;
+ o_orderkey[i] = (index_of_run * 32 + index_in_run + 1);
+ ARROW_DCHECK(1 <= o_orderkey[i] && o_orderkey[i] <= 4 * orders_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CUSTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_CUSTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CUSTKEY));
+
+ // Spec says it must be a random number between 1 and SF*150000 that is not
+ // divisible by 3. Rather than repeatedly generating numbers until we get to
+ // a non-divisible-by-3 number, we just generate a number between
+ // 0 and SF * 50000 - 1, multiply by 3, and then add either 1 or 2.
+ int32_t sf_50k = static_cast<int32_t>(scale_factor_ * 50000);
+ std::uniform_int_distribution<int32_t> base_dist(0, sf_50k - 1);
+ std::uniform_int_distribution<int32_t> offset_dist(1, 2);
+ int32_t* o_custkey = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_CUSTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++)
+ o_custkey[i] = 3 * base_dist(tld.rng) + offset_dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERSTATUS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERSTATUS].kind() == Datum::NONE) {
+ RETURN_NOT_OK(L_LINESTATUS(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERSTATUS));
+
+ char* o_orderstatus = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_ORDERSTATUS].array()->buffers[1]->mutable_data());
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ bool all_f = true;
+ bool all_o = true;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ const char* l_linestatus = reinterpret_cast<const char*>(
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->buffers[1]->data());
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++, batch_offset++) {
+ all_f &= l_linestatus[batch_offset] == 'F';
+ all_o &= l_linestatus[batch_offset] == 'O';
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ ARROW_DCHECK(!(all_f && all_o));
+ if (all_f)
+ o_orderstatus[iorder] = 'F';
+ else if (all_o)
+ o_orderstatus[iorder] = 'O';
+ else
+ o_orderstatus[iorder] = 'P';
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_TOTALPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_TOTALPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(L_EXTENDEDPRICE(thread_index));
+ RETURN_NOT_OK(L_TAX(thread_index));
+ RETURN_NOT_OK(L_DISCOUNT(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_TOTALPRICE));
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ int64_t sum = 0;
+ Decimal128* o_totalprice = reinterpret_cast<Decimal128*>(
+ tld.orders[ORDERS::O_TOTALPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ const Decimal128* l_extendedprice = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->buffers[1]->data());
+ const Decimal128* l_tax = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->data());
+ const Decimal128* l_discount = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->buffers[1]->data());
+
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++, batch_offset++) {
+ int64_t eprice = static_cast<int64_t>(l_extendedprice[batch_offset]);
+ int64_t tax = static_cast<int64_t>(l_tax[batch_offset]);
+ int64_t discount = static_cast<int64_t>(l_discount[batch_offset]);
+ sum += (eprice * (100 + tax) * (100 - discount));
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ sum /= 100 * 100;
+ o_totalprice[iorder] = {sum};
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERDATE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERDATE));
+
+ std::uniform_int_distribution<uint32_t> dist(kStartDate, kEndDate - 151);
+ uint32_t* o_orderdate = reinterpret_cast<uint32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) o_orderdate[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERPRIORITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERPRIORITY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERPRIORITY));
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_ORDERPRIORITY]);
+ std::uniform_int_distribution<int32_t> dist(0, kNumPriorities - 1);
+ char* o_orderpriority = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_ORDERPRIORITY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ const char* str = Priorities[dist(tld.rng)];
+ std::strncpy(o_orderpriority + i * byte_width, str, byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CLERK(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_CLERK].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CLERK));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_CLERK]);
+ int64_t max_clerk_id = static_cast<int64_t>(scale_factor_ * 1000);
+ std::uniform_int_distribution<int64_t> dist(1, max_clerk_id);
+ char* o_clerk = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_CLERK].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ const char* clerk = "Clerk#";
+ const size_t clerk_length = std::strlen(clerk);
+ int64_t clerk_number = dist(tld.rng);
+ char* output = o_clerk + i * byte_width;
+ std::strncpy(output, clerk, byte_width);
+ AppendNumberPaddedToNineDigits(output + clerk_length, clerk_number);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_SHIPPRIORITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_SHIPPRIORITY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_SHIPPRIORITY));
+ int32_t* o_shippriority = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_SHIPPRIORITY].array()->buffers[1]->mutable_data());
+ std::memset(o_shippriority, 0, tld.orders_to_generate * sizeof(int32_t));
+ }
+ return Status::OK();
+ }
+
+ Status O_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.orders[ORDERS::O_COMMENT],
+ g_text.GenerateComments(tld.orders_to_generate, 19, 78, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ Status GenerateRowCounts(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ std::uniform_int_distribution<int> length_dist(1, 7);
+ tld.lineitem_to_generate = 0;
+ tld.items_per_order.clear();
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ int length = length_dist(tld.rng);
+ tld.items_per_order.push_back(length);
+ tld.lineitem_to_generate += length;
+ }
+ int64_t num_batches =
+ (tld.first_batch_offset + tld.lineitem_to_generate + batch_size_ - 1) /
+ batch_size_;
+ tld.lineitem.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.lineitem) {
+ batch.resize(LINEITEM::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocateLineItemBufferIfNeeded(size_t thread_index, size_t ibatch, int column,
+ size_t& out_batch_offset) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.lineitem[ibatch][column].kind() == Datum::NONE) {
+ int32_t byte_width = arrow::internal::GetByteWidth(*kLineitemTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kLineitemTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.lineitem[ibatch][column] = std::move(ad);
+ out_batch_offset = 0;
+ }
+ if (ibatch == 0) out_batch_offset = tld.first_batch_offset;
+
+ return Status::OK();
+ }
+
+ Status L_ORDERKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_ORDERKEY]) {
+ tld.generated_lineitem[LINEITEM::L_ORDERKEY] = true;
+ RETURN_NOT_OK(O_ORDERKEY(thread_index));
+ const int32_t* o_orderkey = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_ORDERKEY, batch_offset));
+ int32_t* l_linenumber =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_ORDERKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_linenumber[batch_offset++] = o_orderkey[iorder];
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_ORDERKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_PARTKEY]) {
+ tld.generated_lineitem[LINEITEM::L_PARTKEY] = true;
+
+ size_t ibatch = 0;
+ int32_t max_partkey = static_cast<int32_t>(scale_factor_ * 200000);
+ std::uniform_int_distribution<int32_t> dist(1, max_partkey);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_PARTKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_partkey =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_partkey[batch_offset] = dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SUPPKEY]) {
+ tld.generated_lineitem[LINEITEM::L_SUPPKEY] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 3);
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset = 0;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SUPPKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_suppkey =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const int32_t* l_partkey = reinterpret_cast<const int32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ for (int64_t i = 0; i < next_run; i++) {
+ int32_t supplier = dist(tld.rng);
+ int32_t partkey = l_partkey[batch_offset];
+ // Fun fact: the parentheses for this expression are unbalanced in the TPC-H
+ // spec.
+ l_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SUPPKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINENUMBER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_LINENUMBER]) {
+ tld.generated_lineitem[LINEITEM::L_LINENUMBER] = true;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_LINENUMBER, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_linenumber =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_LINENUMBER]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++) {
+ l_linenumber[batch_offset++] = (iline + 1);
+ ARROW_DCHECK(1 <= (iline + 1) && (iline + 1) <= 7);
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINENUMBER].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_QUANTITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_QUANTITY]) {
+ tld.generated_lineitem[LINEITEM::L_QUANTITY] = true;
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int64_t> dist(1, 50);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_QUANTITY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_quantity =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_QUANTITY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t i = 0; i < next_run; i++) {
+ // Multiply by 100 because the type is decimal(12, 2), so the decimal goes after
+ // two digits
+ int64_t quantity = dist(tld.rng) * 100;
+ l_quantity[batch_offset++] = {quantity};
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_EXTENDEDPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE]) {
+ tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+ RETURN_NOT_OK(L_QUANTITY(thread_index));
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_EXTENDEDPRICE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ const int32_t* l_partkey = reinterpret_cast<const int32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ const Decimal128* l_quantity = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->buffers[1]->data());
+ Decimal128* l_extendedprice =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ int64_t partkey = static_cast<int64_t>(l_partkey[batch_offset]);
+ // Divide by 100 to recover the integer representation (not Decimal).
+ int64_t quantity = static_cast<int64_t>(l_quantity[batch_offset]) / 100;
+
+ // Spec says to divide by 100, but that happens automatically due to this being
+ // stored to two decimal points.
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ int64_t extended_price = retail_price * quantity;
+ l_extendedprice[batch_offset] = {extended_price};
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_DISCOUNT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_DISCOUNT]) {
+ tld.generated_lineitem[LINEITEM::L_DISCOUNT] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 10);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_DISCOUNT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_discount =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_DISCOUNT]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_discount[batch_offset] = {dist(tld.rng)};
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_TAX(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_TAX]) {
+ tld.generated_lineitem[LINEITEM::L_TAX] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 8);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_TAX, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_tax = reinterpret_cast<Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_tax[batch_offset] = {dist(tld.rng)};
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RETURNFLAG(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_RETURNFLAG]) {
+ tld.generated_lineitem[LINEITEM::L_RETURNFLAG] = true;
+ RETURN_NOT_OK(L_RECEIPTDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<uint32_t> dist;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_RETURNFLAG, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_returnflag =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_receiptdate = reinterpret_cast<const uint32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ if (l_receiptdate[batch_offset] <= kCurrentDate) {
+ uint32_t r = dist(tld.rng);
+ l_returnflag[batch_offset] = (r % 2 == 1) ? 'R' : 'A';
+ } else {
+ l_returnflag[batch_offset] = 'N';
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINESTATUS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_LINESTATUS]) {
+ tld.generated_lineitem[LINEITEM::L_LINESTATUS] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_LINESTATUS, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_linestatus =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_LINESTATUS]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_shipdate =
+ reinterpret_cast<const uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ if (l_shipdate[batch_offset] > kCurrentDate)
+ l_linestatus[batch_offset] = 'O';
+ else
+ l_linestatus[batch_offset] = 'F';
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPDATE]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPDATE] = true;
+ RETURN_NOT_OK(O_ORDERDATE(thread_index));
+ const int32_t* o_orderdate = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(1, 121);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SHIPDATE, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_shipdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_shipdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMITDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_COMMITDATE]) {
+ tld.generated_lineitem[LINEITEM::L_COMMITDATE] = true;
+ const int32_t* o_orderdate = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(30, 90);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_COMMITDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_commitdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_COMMITDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_commitdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMITDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RECEIPTDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_RECEIPTDATE]) {
+ tld.generated_lineitem[LINEITEM::L_RECEIPTDATE] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(1, 30);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_RECEIPTDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_receiptdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_shipdate =
+ reinterpret_cast<const uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_receiptdate[batch_offset] = l_shipdate[batch_offset] + dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPINSTRUCT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT] = true;
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPINSTRUCT]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumInstructions - 1);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_SHIPINSTRUCT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_shipinstruct =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ const char* str = Instructions[dist(tld.rng)];
+ // Note that we don't have to memset the buffer to 0 because strncpy pads each
+ // string with 0's anyway
+ std::strncpy(l_shipinstruct + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPMODE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPMODE]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPMODE] = true;
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPMODE]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumModes - 1);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SHIPMODE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_shipmode =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_SHIPMODE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ const char* str = Modes[dist(tld.rng)];
+ std::strncpy(l_shipmode + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPMODE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_COMMENT]) {
+ tld.generated_lineitem[LINEITEM::L_COMMENT] = true;
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ // Comments are kind of sneaky: we always generate the full batch and then just
+ // bump the length
+ if (tld.lineitem[ibatch][LINEITEM::L_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.lineitem[ibatch][LINEITEM::L_COMMENT],
+ g_text.GenerateComments(batch_size_, 10, 43, tld.rng));
+ batch_offset = 0;
+ }
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ batch_offset += next_run;
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMENT].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ std::vector<Datum> orders;
+ int64_t orders_to_generate;
+ int64_t orderkey_start;
+
+ std::vector<std::vector<Datum>> lineitem;
+ std::vector<int> items_per_order;
+ int64_t lineitem_to_generate;
+ int64_t first_batch_offset;
+ std::bitset<LINEITEM::kNumCols> generated_lineitem;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex orders_output_queue_mutex_;
+ std::mutex lineitem_output_queue_mutex_;
+ std::queue<ExecBatch> orders_output_queue_;
+ std::queue<ExecBatch> lineitem_output_queue_;
+ int64_t batch_size_;
+ float scale_factor_;
+ int64_t orders_rows_to_generate_;
+ int64_t orders_rows_generated_;
+ std::vector<int> orders_cols_;
+ std::vector<int> lineitem_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<size_t> orders_batches_generated_ = {0};
+ std::atomic<size_t> lineitem_batches_generated_ = {0};
+};
+
+class SupplierGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 10000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, gen_list_));
+
+ random::pcg32_fast rng;
+ std::uniform_int_distribution<int64_t> dist(0, rows_to_generate_ - 1);
+ size_t num_special_rows = static_cast<size_t>(5 * scale_factor_);
+ std::unordered_set<int64_t> good_rows_set;
+ while (good_rows_set.size() < num_special_rows) {
+ int64_t row = dist(rng);
+ good_rows_set.insert(row);
+ }
+ std::unordered_set<int64_t> bad_rows_set;
+ while (bad_rows_set.size() < num_special_rows) {
+ int64_t bad_row;
+ do {
+ bad_row = dist(rng);
+ } while (good_rows_set.find(bad_row) != good_rows_set.end());
+ bad_rows_set.insert(bad_row);
+ }
+ good_rows_.clear();
+ bad_rows_.clear();
+ good_rows_.insert(good_rows_.end(), good_rows_set.begin(), good_rows_set.end());
+ bad_rows_.insert(bad_rows_.end(), bad_rows_set.begin(), bad_rows_set.end());
+ std::sort(good_rows_.begin(), good_rows_.end());
+ std::sort(bad_rows_.begin(), bad_rows_.end());
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) tld.rng.seed(seq);
+
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(S_SUPPKEY) \
+ F(S_NAME) \
+ F(S_ADDRESS) \
+ F(S_NATIONKEY) \
+ F(S_PHONE) \
+ F(S_ACCTBAL) \
+ F(S_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct SUPPLIER {
+ enum {
+ FOR_EACH_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) {#col, SUPPLIER::col},
+ const std::unordered_map<std::string, int> kNameMap = {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kGenerators = {FOR_EACH_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(), fixed_size_binary(25), utf8(),
+ int32(), fixed_size_binary(15), decimal(12, 2),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index) {
+ if (done_.load()) return Status::OK();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.suppkey_start = rows_generated_.fetch_add(batch_size_);
+ if (tld.suppkey_start >= rows_to_generate_) return Status::OK();
+
+ tld.to_generate = std::min(batch_size_, rows_to_generate_ - tld.suppkey_start);
+
+ tld.batch.resize(SUPPLIER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for (int col : gen_list_) RETURN_NOT_OK(kGenerators[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for (size_t i = 0; i < gen_list_.size(); i++) {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_outputted_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_outputted_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if (is_last_batch) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ return Status::OK();
+ }
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(kTypes[column], tld.to_generate, {nullptr, std::move(buff)});
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status S_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_SUPPKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_SUPPKEY));
+ int32_t* s_suppkey = reinterpret_cast<int32_t*>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ s_suppkey[irow] = static_cast<int32_t>(tld.suppkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_NAME].kind() == Datum::NONE) {
+ RETURN_NOT_OK(S_SUPPKEY(thread_index));
+ const int32_t* s_suppkey = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->data());
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NAME));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[SUPPLIER::S_NAME]);
+ char* s_name = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_NAME].array()->buffers[1]->mutable_data());
+ // Look man, I'm just following the spec ok? Section 4.2.3 as of March 1 2022
+ const char* supplier = "Supplie#r";
+ const size_t supplier_length = std::strlen(supplier);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = s_name + byte_width * irow;
+ std::strncpy(out, supplier, byte_width);
+ AppendNumberPaddedToNineDigits(out + supplier_length, s_suppkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ADDRESS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_ADDRESS].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[SUPPLIER::S_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status S_NATIONKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_NATIONKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t* s_nationkey = reinterpret_cast<int32_t*>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status S_PHONE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_PHONE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(S_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[SUPPLIER::S_PHONE]);
+ const int32_t* s_nationkey = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->data());
+ char* s_phone = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_PHONE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ GeneratePhoneNumber(s_phone + irow * byte_width, tld.rng, s_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ACCTBAL(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_ACCTBAL].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_ACCTBAL));
+ Decimal128* s_acctbal = reinterpret_cast<Decimal128*>(
+ tld.batch[SUPPLIER::S_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_acctbal[irow] = {dist(tld.rng)};
+ }
+ return Status::OK();
+ }
+
+ Status S_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[SUPPLIER::S_COMMENT],
+ g_text.GenerateComments(tld.to_generate, 25, 100, tld.rng));
+ ModifyComments(thread_index, "Recommends", good_rows_);
+ ModifyComments(thread_index, "Complaints", bad_rows_);
+ }
+ return Status::OK();
+ }
+
+ void ModifyComments(size_t thread_index, const char* review,
+ const std::vector<int64_t>& indices) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ const int32_t* offsets = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[1]->data());
+ char* str = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[2]->mutable_data());
+ const char* customer = "Customer";
+ const int32_t customer_length = static_cast<int32_t>(std::strlen(customer));
+ const int32_t review_length = static_cast<int32_t>(std::strlen(review));
+
+ auto it = std::lower_bound(indices.begin(), indices.end(), tld.suppkey_start);
+ for (; it != indices.end() && *it < tld.suppkey_start + tld.to_generate; it++) {
+ int64_t idx_in_batch = *it - tld.suppkey_start;
+ char* out = str + offsets[idx_in_batch];
+ int32_t str_length = offsets[idx_in_batch + 1] - offsets[idx_in_batch];
+ std::uniform_int_distribution<int32_t> gap_dist(
+ 0, str_length - customer_length - review_length);
+ int32_t gap = gap_dist(tld.rng);
+ int32_t total_length = customer_length + gap + review_length;
+ std::uniform_int_distribution<int32_t> start_dist(0, str_length - total_length);
+ int32_t start = start_dist(tld.rng);
+ std::memcpy(out + start, customer, customer_length);
+ std::memcpy(out + start + customer_length + gap, review, review_length);
+ }
+ }
+
+ struct ThreadLocalData {
+ random::pcg32_fast rng;
+ int64_t suppkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+ std::vector<int64_t> good_rows_;
+ std::vector<int64_t> bad_rows_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_;
+ std::atomic<int64_t> rows_generated_;
+ float scale_factor_;
+ int64_t batch_size_;
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class PartGenerator : public TpchTableGenerator {
+ public:
+ explicit PartGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetPartOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch, gen_->NextPartBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->part_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class PartSuppGenerator : public TpchTableGenerator {
+ public:
+ explicit PartSuppGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetPartSuppOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextPartSuppBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->partsupp_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class CustomerGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, gen_list_));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) tld.rng.seed(seq);
+
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(C_CUSTKEY) \
+ F(C_NAME) \
+ F(C_ADDRESS) \
+ F(C_NATIONKEY) \
+ F(C_PHONE) \
+ F(C_ACCTBAL) \
+ F(C_MKTSEGMENT) \
+ F(C_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct CUSTOMER {
+ enum {
+ FOR_EACH_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) {#col, CUSTOMER::col},
+ const std::unordered_map<std::string, int> kNameMap = {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kGenerators = {FOR_EACH_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(),
+ utf8(),
+ utf8(),
+ int32(),
+ fixed_size_binary(15),
+ decimal(12, 2),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index) {
+ if (done_.load()) return Status::OK();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.custkey_start = rows_generated_.fetch_add(batch_size_);
+ if (tld.custkey_start >= rows_to_generate_) return Status::OK();
+
+ tld.to_generate = std::min(batch_size_, rows_to_generate_ - tld.custkey_start);
+
+ tld.batch.resize(CUSTOMER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for (int col : gen_list_) RETURN_NOT_OK(kGenerators[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for (size_t i = 0; i < gen_list_.size(); i++) {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_generated_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_generated_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if (is_last_batch) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ return Status::OK();
+ }
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(kTypes[column], tld.to_generate, {nullptr, std::move(buff)});
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status C_CUSTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_CUSTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_CUSTKEY));
+ int32_t* c_custkey = reinterpret_cast<int32_t*>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ c_custkey[irow] = static_cast<int32_t>(tld.custkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_NAME].kind() == Datum::NONE) {
+ RETURN_NOT_OK(C_CUSTKEY(thread_index));
+ const int32_t* c_custkey = reinterpret_cast<const int32_t*>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->data());
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ const char* customer = "Customer#";
+ const size_t customer_length = std::strlen(customer);
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ int num_digits = GetNumDigits(c_custkey[irow]);
+ int num_chars = std::max(num_digits, 9);
+ offsets[irow + 1] =
+ static_cast<int32_t>(offsets[irow] + num_chars + customer_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[tld.to_generate]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = str + offsets[irow];
+ std::memcpy(out, customer, customer_length);
+ AppendNumberPaddedToNineDigits(out + customer_length, c_custkey[irow]);
+ }
+ ArrayData ad(utf8(), tld.to_generate,
+ {nullptr, std::move(offset_buff), std::move(str_buff)});
+ tld.batch[CUSTOMER::C_NAME] = std::move(ad);
+ }
+ return Status::OK();
+ }
+
+ Status C_ADDRESS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_ADDRESS].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[CUSTOMER::C_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status C_NATIONKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_NATIONKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t* c_nationkey = reinterpret_cast<int32_t*>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status C_PHONE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_PHONE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(C_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[CUSTOMER::C_PHONE]);
+ const int32_t* c_nationkey = reinterpret_cast<const int32_t*>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->data());
+ char* c_phone = reinterpret_cast<char*>(
+ tld.batch[CUSTOMER::C_PHONE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ GeneratePhoneNumber(c_phone + irow * byte_width, tld.rng, c_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_ACCTBAL(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_ACCTBAL].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_ACCTBAL));
+ Decimal128* c_acctbal = reinterpret_cast<Decimal128*>(
+ tld.batch[CUSTOMER::C_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_acctbal[irow] = {dist(tld.rng)};
+ }
+ return Status::OK();
+ }
+
+ Status C_MKTSEGMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_MKTSEGMENT].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_MKTSEGMENT));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[CUSTOMER::C_MKTSEGMENT]);
+ char* c_mktsegment = reinterpret_cast<char*>(
+ tld.batch[CUSTOMER::C_MKTSEGMENT].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int32_t> dist(0, kNumSegments - 1);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = c_mktsegment + irow * byte_width;
+ int str_idx = dist(tld.rng);
+ std::strncpy(out, Segments[str_idx], byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[CUSTOMER::C_COMMENT],
+ g_text.GenerateComments(tld.to_generate, 29, 116, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ random::pcg32_fast rng;
+ int64_t custkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_{0};
+ std::atomic<int64_t> rows_generated_ = {0};
+ float scale_factor_{0};
+ int64_t batch_size_{0};
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class OrdersGenerator : public TpchTableGenerator {
+ public:
+ explicit OrdersGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetOrdersOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch, gen_->NextOrdersBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->orders_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class LineitemGenerator : public TpchTableGenerator {
+ public:
+ explicit LineitemGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetLineItemOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextLineItemBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->lineitem_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class NationGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, float /*scale_factor*/,
+ int64_t /*batch_size*/) override {
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, column_indices_));
+ rng_.seed(arrow_vendored::pcg_extras::seed_seq_from<std::random_device>{});
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t /*num_threads*/, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback /*schedule_task_callback*/) override {
+ std::shared_ptr<Buffer> N_NATIONKEY_buffer =
+ Buffer::Wrap(kNationKey, sizeof(kNationKey));
+ ArrayData N_NATIONKEY_arraydata(int32(), kRowCount,
+ {nullptr, std::move(N_NATIONKEY_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> N_NAME_buffer,
+ AllocateBuffer(kRowCount * kNameByteWidth));
+ char* N_NAME = reinterpret_cast<char*>(N_NAME_buffer->mutable_data());
+ for (size_t i = 0; i < kRowCount; i++)
+ std::strncpy(N_NAME + kNameByteWidth * i, kCountryNames[i], kNameByteWidth);
+ ArrayData N_NAME_arraydata(fixed_size_binary(kNameByteWidth), kRowCount,
+ {nullptr, std::move(N_NAME_buffer)});
+
+ std::shared_ptr<Buffer> N_REGIONKEY_buffer =
+ Buffer::Wrap(kRegionKey, sizeof(kRegionKey));
+ ArrayData N_REGIONKEY_arraydata(int32(), kRowCount,
+ {nullptr, std::move(N_REGIONKEY_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(Datum N_COMMENT_datum,
+ g_text.GenerateComments(kRowCount, 31, 114, rng_));
+
+ std::vector<Datum> fields = {
+ std::move(N_NATIONKEY_arraydata), std::move(N_NAME_arraydata),
+ std::move(N_REGIONKEY_arraydata), std::move(N_COMMENT_datum)};
+
+ std::vector<Datum> result;
+ for (const int& col : column_indices_) result.push_back(fields[col]);
+ ARROW_ASSIGN_OR_RAISE(ExecBatch batch, ExecBatch::Make(std::move(result)));
+ output_callback(std::move(batch));
+ finished_callback(static_cast<int64_t>(1));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ random::pcg32_fast rng_;
+
+ static constexpr size_t kRowCount = 25;
+ static constexpr int32_t kNameByteWidth = 25;
+ const int32_t kNationKey[kRowCount] = {0, 1, 2, 3, 4, 5, 6, 7, 8,
+ 9, 10, 11, 12, 13, 14, 15, 16, 17,
+ 18, 19, 20, 21, 22, 23, 24};
+ const char* kCountryNames[kRowCount] = {
+ "ALGERIA", "ARGENTINA", "BRAZIL", "CANADA", "EGYPT",
+ "ETHIOPIA", "FRANCE", "GERMANY", "INDIA", "INDONESIA",
+ "IRAN", "IRAQ", "JAPAN", "JORDAN", "KENYA",
+ "MOROCCO", "MOZAMBIQUE", "PERU", "CHINA", "ROMANIA",
+ "SAUDI ARABIA", "VIETNAM", "RUSSIA", "UNITED KINGDOM", "UNITED STATES"};
+ const int32_t kRegionKey[kRowCount] = {0, 1, 1, 1, 4, 0, 3, 3, 2, 2, 4, 4, 2,
+ 4, 0, 0, 0, 1, 2, 3, 4, 2, 3, 3, 1};
+
+ struct NATION {
+ enum {
Review comment:
You can instead write:
```suggestion
enum class NATION {
```
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
Review comment:
Can you instead use e.g. a `std::vector<std::string>`? We're trying to write idiomatic C++ code as much as possible.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
Review comment:
Just for the record, TPC-H specifies that phone numbers never have a leading zero?
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
Review comment:
I may be mistaken (this code is difficult to follow and you didn't provide any comments), but it looks like this is doing home-grown thread synchronization. Is there a strong reason not to use the existing `ThreadPool` facility?
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
Review comment:
Again, this could probably use one of the formatting facilities.
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
Review comment:
Also, note you can use a verification class/struct that would expose the VerifyXXX methods and accumulate statistics about the data in member attributes. This would probably be more idiomatic.
##########
File path: cpp/src/arrow/compute/kernels/vector_sort.cc
##########
@@ -893,7 +893,8 @@ class TableSorter {
TableSorter(ExecContext* ctx, uint64_t* indices_begin, uint64_t* indices_end,
const Table& table, const SortOptions& options)
- : ctx_(ctx),
Review comment:
@save-buffer Do you remember where the "resetting the error" part happened?
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
Review comment:
It would be better to call `GetRandomSeed` (see `arrow/util/io_util.h`).
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ if (!partsupp_output_queue_.empty()) {
+ ExecBatch result = std::move(partsupp_output_queue_.front());
+ partsupp_output_queue_.pop();
+ return std::move(result);
+ }
+ }
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+ if (!part_cols_.empty()) {
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch part_batch,
+ ExecBatch::Make(std::move(part_result)));
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ part_output_queue_.emplace(std::move(part_batch));
+ }
+ }
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ for (size_t i = 1; i < partsupp_results.size(); i++)
+ partsupp_output_queue_.emplace(std::move(partsupp_results[i]));
+ }
+ return std::move(partsupp_results[0]);
+ }
+
+ private:
+#define FOR_EACH_PART_COLUMN(F) \
+ F(P_PARTKEY) \
+ F(P_NAME) \
+ F(P_MFGR) \
+ F(P_BRAND) \
+ F(P_TYPE) \
+ F(P_SIZE) \
+ F(P_CONTAINER) \
+ F(P_RETAILPRICE) \
+ F(P_COMMENT)
+
+#define FOR_EACH_PARTSUPP_COLUMN(F) \
+ F(PS_PARTKEY) \
+ F(PS_SUPPKEY) \
+ F(PS_AVAILQTY) \
+ F(PS_SUPPLYCOST) \
+ F(PS_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct PART {
+ enum {
+ FOR_EACH_PART_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct PARTSUPP {
+ enum {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, PART::col},
+ const std::unordered_map<std::string, int> kPartNameMap = {
+ FOR_EACH_PART_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, PARTSUPP::col},
+ const std::unordered_map<std::string, int> kPartsuppNameMap = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kPartGenerators = {FOR_EACH_PART_COLUMN(MAKE_FN_ARRAY)};
+ std::vector<GenerateColumnFn> kPartsuppGenerators = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kPartTypes = {
+ int32(), utf8(), fixed_size_binary(25), fixed_size_binary(10),
+ utf8(), int32(), fixed_size_binary(10), decimal(12, 2),
+ utf8(),
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kPartsuppTypes = {
+ int32(), int32(), int32(), decimal(12, 2), utf8(),
+ };
+
+ Status AllocatePartBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.part[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.part_to_generate * byte_width));
+ ArrayData ad(kPartTypes[column], tld.part_to_generate, {nullptr, std::move(buff)});
+ tld.part[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status P_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_PARTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_PARTKEY));
+ int32_t* p_partkey = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) {
+ p_partkey[i] = static_cast<int32_t>(tld.partkey_start + i + 1);
+ ARROW_DCHECK(1 <= p_partkey[i] && p_partkey[i] <= part_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_NAME].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(
+ 0, static_cast<uint8_t>(kNumNameParts - 1));
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dist(tld.rng));
+ tld.string_indices[irow * 5 + ipart] = name_part_index;
+ string_length += std::strlen(NameParts[name_part_index]);
+ }
+ // Add 4 because there is a space between each word (i.e. four spaces)
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length + 4);
+ }
+ // Add an extra byte for the space after in the very last string.
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate] + 1));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 5 + ipart];
+ const char* part = NameParts[name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ *row++ = ' ';
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_NAME], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_NAME] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_MFGR(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_MFGR].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* manufacturer = "Manufacturer#";
+ const size_t manufacturer_length = std::strlen(manufacturer);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_MFGR));
+ char* p_mfgr = reinterpret_cast<char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ std::strncpy(p_mfgr + byte_width * irow, manufacturer, byte_width);
+ char mfgr_id = '0' + dist(tld.rng);
+ *(p_mfgr + byte_width * irow + manufacturer_length) = mfgr_id;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_BRAND(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_BRAND].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_MFGR(thread_index));
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* brand = "Brand#";
+ const size_t brand_length = std::strlen(brand);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_BRAND));
+ const char* p_mfgr = reinterpret_cast<const char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->data());
+ char* p_brand = reinterpret_cast<char*>(
+ tld.part[PART::P_BRAND].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_BRAND]);
+ int32_t mfgr_byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ const size_t mfgr_id_offset = std::strlen("Manufacturer#");
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = p_brand + byte_width * irow;
+ char mfgr_id = *(p_mfgr + irow * mfgr_byte_width + mfgr_id_offset);
+ char brand_id = '0' + dist(tld.rng);
+ std::strncpy(row, brand, byte_width);
+ *(row + brand_length) = mfgr_id;
+ *(row + brand_length + 1) = brand_id;
+ irow += 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_TYPE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_TYPE].kind() == Datum::NONE) {
+ using D = std::uniform_int_distribution<uint32_t>;
+ D dists[] = {
+ D{0, static_cast<uint8_t>(kNumTypes_1 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_2 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_3 - 1)},
+ };
+
+ const char** types[] = {Types_1, Types_2, Types_3};
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dists[ipart](tld.rng));
+ tld.string_indices[irow * 3 + ipart] = name_part_index;
+ string_length += std::strlen(types[ipart][name_part_index]);
+ }
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate]));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 3 + ipart];
+ const char* part = types[ipart][name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_TYPE], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_TYPE] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_SIZE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_SIZE].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(1, 50);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_SIZE));
+ int32_t* p_size = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_SIZE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) p_size[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status P_CONTAINER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_CONTAINER].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist1(
+ 0, static_cast<uint8_t>(kNumContainers_1 - 1));
+ std::uniform_int_distribution<int> dist2(
+ 0, static_cast<uint8_t>(kNumContainers_2 - 1));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_CONTAINER));
+ char* p_container = reinterpret_cast<char*>(
+ tld.part[PART::P_CONTAINER].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_CONTAINER]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int container1_idx = dist1(tld.rng);
+ int container2_idx = dist2(tld.rng);
+ const char* container1 = Containers_1[container1_idx];
+ const char* container2 = Containers_2[container2_idx];
+ size_t container1_length = std::strlen(container1);
+ size_t container2_length = std::strlen(container2);
+
+ char* row = p_container + byte_width * irow;
+ std::strncpy(row, container1, byte_width);
+ std::memcpy(row + container1_length, container2, container2_length);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_RETAILPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_RETAILPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_RETAILPRICE));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+ Decimal128* p_retailprice = reinterpret_cast<Decimal128*>(
+ tld.part[PART::P_RETAILPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int32_t partkey = p_partkey[irow];
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ p_retailprice[irow] = {retail_price};
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.part[PART::P_COMMENT],
+ g_text.GenerateComments(tld.part_to_generate, 5, 22, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ int64_t PartsuppBatchesToGenerate(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t num_batches = (ps_to_generate + batch_size_ - 1) / batch_size_;
+ return num_batches;
+ }
+
+ Status InitPartsupp(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.generated_partsupp.reset();
+ int64_t num_batches = PartsuppBatchesToGenerate(thread_index);
+ tld.partsupp.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.partsupp) {
+ batch.resize(PARTSUPP::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocatePartSuppBatch(size_t thread_index, size_t ibatch, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartsuppTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kPartsuppTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.partsupp[ibatch][column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status PS_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_PARTKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_PARTKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_PARTKEY));
+ int32_t* ps_partkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run; ipartsupp++, irun++)
+ ps_partkey[batch_offset++] = p_partkey[ipart];
+
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPKEY));
+ int32_t* ps_suppkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run;
+ ipartsupp++, irun++) {
+ int32_t supplier = static_cast<int32_t>(ipartsupp);
+ int32_t partkey = p_partkey[ipart];
+ ps_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
Review comment:
What is this? Some magic formula from TPC-H? Can you please add a comment?
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ if (!partsupp_output_queue_.empty()) {
+ ExecBatch result = std::move(partsupp_output_queue_.front());
+ partsupp_output_queue_.pop();
+ return std::move(result);
+ }
+ }
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+ if (!part_cols_.empty()) {
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch part_batch,
+ ExecBatch::Make(std::move(part_result)));
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ part_output_queue_.emplace(std::move(part_batch));
+ }
+ }
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ for (size_t i = 1; i < partsupp_results.size(); i++)
+ partsupp_output_queue_.emplace(std::move(partsupp_results[i]));
+ }
+ return std::move(partsupp_results[0]);
+ }
+
+ private:
+#define FOR_EACH_PART_COLUMN(F) \
+ F(P_PARTKEY) \
+ F(P_NAME) \
+ F(P_MFGR) \
+ F(P_BRAND) \
+ F(P_TYPE) \
+ F(P_SIZE) \
+ F(P_CONTAINER) \
+ F(P_RETAILPRICE) \
+ F(P_COMMENT)
+
+#define FOR_EACH_PARTSUPP_COLUMN(F) \
+ F(PS_PARTKEY) \
+ F(PS_SUPPKEY) \
+ F(PS_AVAILQTY) \
+ F(PS_SUPPLYCOST) \
+ F(PS_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct PART {
+ enum {
+ FOR_EACH_PART_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct PARTSUPP {
+ enum {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, PART::col},
+ const std::unordered_map<std::string, int> kPartNameMap = {
+ FOR_EACH_PART_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, PARTSUPP::col},
+ const std::unordered_map<std::string, int> kPartsuppNameMap = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kPartGenerators = {FOR_EACH_PART_COLUMN(MAKE_FN_ARRAY)};
+ std::vector<GenerateColumnFn> kPartsuppGenerators = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kPartTypes = {
+ int32(), utf8(), fixed_size_binary(25), fixed_size_binary(10),
+ utf8(), int32(), fixed_size_binary(10), decimal(12, 2),
+ utf8(),
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kPartsuppTypes = {
+ int32(), int32(), int32(), decimal(12, 2), utf8(),
+ };
+
+ Status AllocatePartBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.part[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.part_to_generate * byte_width));
+ ArrayData ad(kPartTypes[column], tld.part_to_generate, {nullptr, std::move(buff)});
+ tld.part[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status P_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_PARTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_PARTKEY));
+ int32_t* p_partkey = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) {
+ p_partkey[i] = static_cast<int32_t>(tld.partkey_start + i + 1);
+ ARROW_DCHECK(1 <= p_partkey[i] && p_partkey[i] <= part_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_NAME].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(
+ 0, static_cast<uint8_t>(kNumNameParts - 1));
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dist(tld.rng));
+ tld.string_indices[irow * 5 + ipart] = name_part_index;
+ string_length += std::strlen(NameParts[name_part_index]);
+ }
+ // Add 4 because there is a space between each word (i.e. four spaces)
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length + 4);
+ }
+ // Add an extra byte for the space after in the very last string.
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate] + 1));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 5 + ipart];
+ const char* part = NameParts[name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ *row++ = ' ';
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_NAME], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_NAME] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_MFGR(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_MFGR].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* manufacturer = "Manufacturer#";
+ const size_t manufacturer_length = std::strlen(manufacturer);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_MFGR));
+ char* p_mfgr = reinterpret_cast<char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ std::strncpy(p_mfgr + byte_width * irow, manufacturer, byte_width);
+ char mfgr_id = '0' + dist(tld.rng);
+ *(p_mfgr + byte_width * irow + manufacturer_length) = mfgr_id;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_BRAND(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_BRAND].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_MFGR(thread_index));
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* brand = "Brand#";
+ const size_t brand_length = std::strlen(brand);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_BRAND));
+ const char* p_mfgr = reinterpret_cast<const char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->data());
+ char* p_brand = reinterpret_cast<char*>(
+ tld.part[PART::P_BRAND].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_BRAND]);
+ int32_t mfgr_byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ const size_t mfgr_id_offset = std::strlen("Manufacturer#");
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = p_brand + byte_width * irow;
+ char mfgr_id = *(p_mfgr + irow * mfgr_byte_width + mfgr_id_offset);
+ char brand_id = '0' + dist(tld.rng);
+ std::strncpy(row, brand, byte_width);
+ *(row + brand_length) = mfgr_id;
+ *(row + brand_length + 1) = brand_id;
+ irow += 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_TYPE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_TYPE].kind() == Datum::NONE) {
+ using D = std::uniform_int_distribution<uint32_t>;
+ D dists[] = {
+ D{0, static_cast<uint8_t>(kNumTypes_1 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_2 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_3 - 1)},
+ };
+
+ const char** types[] = {Types_1, Types_2, Types_3};
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dists[ipart](tld.rng));
+ tld.string_indices[irow * 3 + ipart] = name_part_index;
+ string_length += std::strlen(types[ipart][name_part_index]);
+ }
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate]));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 3 + ipart];
+ const char* part = types[ipart][name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_TYPE], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_TYPE] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_SIZE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_SIZE].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(1, 50);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_SIZE));
+ int32_t* p_size = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_SIZE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) p_size[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status P_CONTAINER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_CONTAINER].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist1(
+ 0, static_cast<uint8_t>(kNumContainers_1 - 1));
+ std::uniform_int_distribution<int> dist2(
+ 0, static_cast<uint8_t>(kNumContainers_2 - 1));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_CONTAINER));
+ char* p_container = reinterpret_cast<char*>(
+ tld.part[PART::P_CONTAINER].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_CONTAINER]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int container1_idx = dist1(tld.rng);
+ int container2_idx = dist2(tld.rng);
+ const char* container1 = Containers_1[container1_idx];
+ const char* container2 = Containers_2[container2_idx];
+ size_t container1_length = std::strlen(container1);
+ size_t container2_length = std::strlen(container2);
+
+ char* row = p_container + byte_width * irow;
+ std::strncpy(row, container1, byte_width);
+ std::memcpy(row + container1_length, container2, container2_length);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_RETAILPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_RETAILPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_RETAILPRICE));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+ Decimal128* p_retailprice = reinterpret_cast<Decimal128*>(
+ tld.part[PART::P_RETAILPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int32_t partkey = p_partkey[irow];
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ p_retailprice[irow] = {retail_price};
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.part[PART::P_COMMENT],
+ g_text.GenerateComments(tld.part_to_generate, 5, 22, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ int64_t PartsuppBatchesToGenerate(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t num_batches = (ps_to_generate + batch_size_ - 1) / batch_size_;
+ return num_batches;
+ }
+
+ Status InitPartsupp(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.generated_partsupp.reset();
+ int64_t num_batches = PartsuppBatchesToGenerate(thread_index);
+ tld.partsupp.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.partsupp) {
+ batch.resize(PARTSUPP::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocatePartSuppBatch(size_t thread_index, size_t ibatch, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartsuppTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kPartsuppTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.partsupp[ibatch][column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status PS_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_PARTKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_PARTKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_PARTKEY));
+ int32_t* ps_partkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run; ipartsupp++, irun++)
+ ps_partkey[batch_offset++] = p_partkey[ipart];
+
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPKEY));
+ int32_t* ps_suppkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run;
+ ipartsupp++, irun++) {
+ int32_t supplier = static_cast<int32_t>(ipartsupp);
+ int32_t partkey = p_partkey[ipart];
+ ps_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_AVAILQTY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_AVAILQTY]) {
+ tld.generated_partsupp[PARTSUPP::PS_AVAILQTY] = true;
+ std::uniform_int_distribution<int32_t> dist(1, 9999);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_AVAILQTY));
+ int32_t* ps_availqty =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for (int64_t irun = 0; irun < next_run; irun++) ps_availqty[irun] = dist(tld.rng);
+
+ tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPLYCOST(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST] = true;
+ std::uniform_int_distribution<int64_t> dist(100, 100000);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(
+ AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPLYCOST));
+ Decimal128* ps_supplycost =
+ reinterpret_cast<Decimal128*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for (int64_t irun = 0; irun < next_run; irun++)
+ ps_supplycost[irun] = {dist(tld.rng)};
+
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PARTSUPP::PS_COMMENT].kind() == Datum::NONE) {
+ int64_t irow = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ int64_t num_rows = std::min(batch_size_, ps_to_generate - irow);
+ ARROW_ASSIGN_OR_RAISE(tld.partsupp[ibatch][PARTSUPP::PS_COMMENT],
+ g_text.GenerateComments(num_rows, 49, 198, tld.rng));
+ irow += num_rows;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ std::vector<Datum> part;
+ std::vector<int8_t> string_indices;
+ int64_t part_to_generate{0};
+ int64_t partkey_start{0};
+
+ std::vector<std::vector<Datum>> partsupp;
+ std::bitset<PARTSUPP::kNumCols> generated_partsupp;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex part_output_queue_mutex_;
+ std::mutex partsupp_output_queue_mutex_;
+ std::queue<ExecBatch> part_output_queue_;
+ std::queue<ExecBatch> partsupp_output_queue_;
+ int64_t batch_size_{0};
+ float scale_factor_{0};
+ int64_t part_rows_to_generate_{0};
+ int64_t part_rows_generated_{0};
+ std::vector<int> part_cols_;
+ std::vector<int> partsupp_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<int64_t> part_batches_generated_ = {0};
+ std::atomic<int64_t> partsupp_batches_generated_ = {0};
+ static constexpr int64_t kPartSuppRowsPerPart = 4;
+};
+
+class OrdersAndLineItemGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ tld.items_per_order.resize(batch_size_);
+ tld.rng.seed(seq);
+ }
+ orders_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000 * 10);
+ }
+ return Status::OK();
+ }
+
+ int64_t orders_batches_generated() const { return orders_batches_generated_.load(); }
+
+ int64_t lineitem_batches_generated() const {
+ return lineitem_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetOrdersOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kOrdersTypes, kOrdersNameMap, orders_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetLineItemOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kLineitemTypes, kLineitemNameMap, lineitem_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextOrdersBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if (!orders_output_queue_.empty()) {
+ ExecBatch batch = std::move(orders_output_queue_.front());
+ orders_output_queue_.pop();
+ return std::move(batch);
+ } else if (orders_rows_generated_ == orders_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate =
+ std::min(batch_size_, orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.first_batch_offset = 0;
+ tld.generated_lineitem.reset();
+
+ for (int col : orders_cols_) RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for (int col : lineitem_cols_) RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for (size_t i = 0; i < orders_cols_.size(); i++) {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ if (!lineitem_cols_.empty()) {
+ std::vector<ExecBatch> lineitem_results;
+ for (size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++) {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for (size_t icol = 0; icol < lineitem_cols_.size(); icol++) {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(lineitem_output_queue_mutex_);
+ for (ExecBatch& eb : lineitem_results) {
+ lineitem_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(orders_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextLineItemBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ExecBatch queued;
+ bool from_queue = false;
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ if (!lineitem_output_queue_.empty()) {
+ queued = std::move(lineitem_output_queue_.front());
+ lineitem_output_queue_.pop();
+ from_queue = true;
+ }
+ }
+ tld.first_batch_offset = 0;
+ if (from_queue) {
+ ARROW_DCHECK(queued.length <= batch_size_);
+ tld.first_batch_offset = queued.length;
+ if (queued.length == batch_size_) return std::move(queued);
+ }
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if (orders_rows_generated_ == orders_rows_to_generate_) {
+ if (from_queue) return std::move(queued);
+ return util::nullopt;
+ }
+
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate =
+ std::min(batch_size_, orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1ll);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.generated_lineitem.reset();
+ if (from_queue) {
+ lineitem_batches_generated_.fetch_sub(1);
+ for (size_t i = 0; i < lineitem_cols_.size(); i++)
+ if (tld.lineitem[0][lineitem_cols_[i]].kind() == Datum::NONE)
+ tld.lineitem[0][lineitem_cols_[i]] = std::move(queued[i]);
+ }
+
+ for (int col : orders_cols_) RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for (int col : lineitem_cols_) RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ if (!orders_cols_.empty()) {
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for (size_t i = 0; i < orders_cols_.size(); i++) {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch orders_batch,
+ ExecBatch::Make(std::move(orders_result)));
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ orders_output_queue_.emplace(std::move(orders_batch));
+ }
+ }
+ std::vector<ExecBatch> lineitem_results;
+ for (size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++) {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for (size_t icol = 0; icol < lineitem_cols_.size(); icol++) {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ lineitem_batches_generated_.fetch_add(static_cast<int64_t>(lineitem_results.size()));
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ for (size_t i = 1; i < lineitem_results.size(); i++)
+ lineitem_output_queue_.emplace(std::move(lineitem_results[i]));
+ }
+ return std::move(lineitem_results[0]);
+ }
+
+ private:
+#define FOR_EACH_ORDERS_COLUMN(F) \
+ F(O_ORDERKEY) \
+ F(O_CUSTKEY) \
+ F(O_ORDERSTATUS) \
+ F(O_TOTALPRICE) \
+ F(O_ORDERDATE) \
+ F(O_ORDERPRIORITY) \
+ F(O_CLERK) \
+ F(O_SHIPPRIORITY) \
+ F(O_COMMENT)
+
+#define FOR_EACH_LINEITEM_COLUMN(F) \
+ F(L_ORDERKEY) \
+ F(L_PARTKEY) \
+ F(L_SUPPKEY) \
+ F(L_LINENUMBER) \
+ F(L_QUANTITY) \
+ F(L_EXTENDEDPRICE) \
+ F(L_DISCOUNT) \
+ F(L_TAX) \
+ F(L_RETURNFLAG) \
+ F(L_LINESTATUS) \
+ F(L_SHIPDATE) \
+ F(L_COMMITDATE) \
+ F(L_RECEIPTDATE) \
+ F(L_SHIPINSTRUCT) \
+ F(L_SHIPMODE) \
+ F(L_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct ORDERS {
+ enum {
+ FOR_EACH_ORDERS_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct LINEITEM {
+ enum {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, ORDERS::col},
+ const std::unordered_map<std::string, int> kOrdersNameMap = {
+ FOR_EACH_ORDERS_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, LINEITEM::col},
+ const std::unordered_map<std::string, int> kLineitemNameMap = {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ const std::vector<GenerateColumnFn> kOrdersGenerators = {
+ FOR_EACH_ORDERS_COLUMN(MAKE_FN_ARRAY)};
+ const std::vector<GenerateColumnFn> kLineitemGenerators = {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kOrdersTypes = {int32(),
+ int32(),
+ fixed_size_binary(1),
+ decimal(12, 2),
+ date32(),
+ fixed_size_binary(15),
+ fixed_size_binary(15),
+ int32(),
+ utf8()};
+
+ const std::vector<std::shared_ptr<DataType>> kLineitemTypes = {
+ int32(),
+ int32(),
+ int32(),
+ int32(),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ fixed_size_binary(1),
+ fixed_size_binary(1),
+ date32(),
+ date32(),
+ date32(),
+ fixed_size_binary(25),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status AllocateOrdersBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.orders[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.orders_to_generate * byte_width));
+ ArrayData ad(kOrdersTypes[column], tld.orders_to_generate,
+ {nullptr, std::move(buff)});
+ tld.orders[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status O_ORDERKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERKEY));
+ int32_t* o_orderkey = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ int32_t orderkey_index = static_cast<int32_t>(tld.orderkey_start + i);
+ int32_t index_of_run = orderkey_index / 8;
+ int32_t index_in_run = orderkey_index % 8;
+ o_orderkey[i] = (index_of_run * 32 + index_in_run + 1);
+ ARROW_DCHECK(1 <= o_orderkey[i] && o_orderkey[i] <= 4 * orders_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CUSTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_CUSTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CUSTKEY));
+
+ // Spec says it must be a random number between 1 and SF*150000 that is not
+ // divisible by 3. Rather than repeatedly generating numbers until we get to
+ // a non-divisible-by-3 number, we just generate a number between
+ // 0 and SF * 50000 - 1, multiply by 3, and then add either 1 or 2.
+ int32_t sf_50k = static_cast<int32_t>(scale_factor_ * 50000);
+ std::uniform_int_distribution<int32_t> base_dist(0, sf_50k - 1);
+ std::uniform_int_distribution<int32_t> offset_dist(1, 2);
+ int32_t* o_custkey = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_CUSTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++)
+ o_custkey[i] = 3 * base_dist(tld.rng) + offset_dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERSTATUS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERSTATUS].kind() == Datum::NONE) {
+ RETURN_NOT_OK(L_LINESTATUS(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERSTATUS));
+
+ char* o_orderstatus = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_ORDERSTATUS].array()->buffers[1]->mutable_data());
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ bool all_f = true;
+ bool all_o = true;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ const char* l_linestatus = reinterpret_cast<const char*>(
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->buffers[1]->data());
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++, batch_offset++) {
+ all_f &= l_linestatus[batch_offset] == 'F';
+ all_o &= l_linestatus[batch_offset] == 'O';
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ ARROW_DCHECK(!(all_f && all_o));
+ if (all_f)
+ o_orderstatus[iorder] = 'F';
+ else if (all_o)
+ o_orderstatus[iorder] = 'O';
+ else
+ o_orderstatus[iorder] = 'P';
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_TOTALPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_TOTALPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(L_EXTENDEDPRICE(thread_index));
+ RETURN_NOT_OK(L_TAX(thread_index));
+ RETURN_NOT_OK(L_DISCOUNT(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_TOTALPRICE));
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ int64_t sum = 0;
+ Decimal128* o_totalprice = reinterpret_cast<Decimal128*>(
+ tld.orders[ORDERS::O_TOTALPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ const Decimal128* l_extendedprice = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->buffers[1]->data());
+ const Decimal128* l_tax = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->data());
+ const Decimal128* l_discount = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->buffers[1]->data());
+
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++, batch_offset++) {
+ int64_t eprice = static_cast<int64_t>(l_extendedprice[batch_offset]);
+ int64_t tax = static_cast<int64_t>(l_tax[batch_offset]);
+ int64_t discount = static_cast<int64_t>(l_discount[batch_offset]);
+ sum += (eprice * (100 + tax) * (100 - discount));
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ sum /= 100 * 100;
+ o_totalprice[iorder] = {sum};
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERDATE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERDATE));
+
+ std::uniform_int_distribution<uint32_t> dist(kStartDate, kEndDate - 151);
+ uint32_t* o_orderdate = reinterpret_cast<uint32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) o_orderdate[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERPRIORITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERPRIORITY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERPRIORITY));
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_ORDERPRIORITY]);
+ std::uniform_int_distribution<int32_t> dist(0, kNumPriorities - 1);
+ char* o_orderpriority = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_ORDERPRIORITY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ const char* str = Priorities[dist(tld.rng)];
+ std::strncpy(o_orderpriority + i * byte_width, str, byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CLERK(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_CLERK].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CLERK));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_CLERK]);
+ int64_t max_clerk_id = static_cast<int64_t>(scale_factor_ * 1000);
+ std::uniform_int_distribution<int64_t> dist(1, max_clerk_id);
+ char* o_clerk = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_CLERK].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ const char* clerk = "Clerk#";
+ const size_t clerk_length = std::strlen(clerk);
+ int64_t clerk_number = dist(tld.rng);
+ char* output = o_clerk + i * byte_width;
+ std::strncpy(output, clerk, byte_width);
+ AppendNumberPaddedToNineDigits(output + clerk_length, clerk_number);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_SHIPPRIORITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_SHIPPRIORITY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_SHIPPRIORITY));
+ int32_t* o_shippriority = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_SHIPPRIORITY].array()->buffers[1]->mutable_data());
+ std::memset(o_shippriority, 0, tld.orders_to_generate * sizeof(int32_t));
+ }
+ return Status::OK();
+ }
+
+ Status O_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.orders[ORDERS::O_COMMENT],
+ g_text.GenerateComments(tld.orders_to_generate, 19, 78, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ Status GenerateRowCounts(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ std::uniform_int_distribution<int> length_dist(1, 7);
+ tld.lineitem_to_generate = 0;
+ tld.items_per_order.clear();
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ int length = length_dist(tld.rng);
+ tld.items_per_order.push_back(length);
+ tld.lineitem_to_generate += length;
+ }
+ int64_t num_batches =
+ (tld.first_batch_offset + tld.lineitem_to_generate + batch_size_ - 1) /
+ batch_size_;
+ tld.lineitem.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.lineitem) {
+ batch.resize(LINEITEM::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocateLineItemBufferIfNeeded(size_t thread_index, size_t ibatch, int column,
+ size_t& out_batch_offset) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.lineitem[ibatch][column].kind() == Datum::NONE) {
+ int32_t byte_width = arrow::internal::GetByteWidth(*kLineitemTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kLineitemTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.lineitem[ibatch][column] = std::move(ad);
+ out_batch_offset = 0;
+ }
+ if (ibatch == 0) out_batch_offset = tld.first_batch_offset;
+
+ return Status::OK();
+ }
+
+ Status L_ORDERKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_ORDERKEY]) {
+ tld.generated_lineitem[LINEITEM::L_ORDERKEY] = true;
+ RETURN_NOT_OK(O_ORDERKEY(thread_index));
+ const int32_t* o_orderkey = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_ORDERKEY, batch_offset));
+ int32_t* l_linenumber =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_ORDERKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_linenumber[batch_offset++] = o_orderkey[iorder];
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_ORDERKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_PARTKEY]) {
+ tld.generated_lineitem[LINEITEM::L_PARTKEY] = true;
+
+ size_t ibatch = 0;
+ int32_t max_partkey = static_cast<int32_t>(scale_factor_ * 200000);
+ std::uniform_int_distribution<int32_t> dist(1, max_partkey);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_PARTKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_partkey =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_partkey[batch_offset] = dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SUPPKEY]) {
+ tld.generated_lineitem[LINEITEM::L_SUPPKEY] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 3);
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset = 0;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SUPPKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_suppkey =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const int32_t* l_partkey = reinterpret_cast<const int32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ for (int64_t i = 0; i < next_run; i++) {
+ int32_t supplier = dist(tld.rng);
+ int32_t partkey = l_partkey[batch_offset];
+ // Fun fact: the parentheses for this expression are unbalanced in the TPC-H
+ // spec.
+ l_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SUPPKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINENUMBER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_LINENUMBER]) {
+ tld.generated_lineitem[LINEITEM::L_LINENUMBER] = true;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_LINENUMBER, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_linenumber =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_LINENUMBER]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++) {
+ l_linenumber[batch_offset++] = (iline + 1);
+ ARROW_DCHECK(1 <= (iline + 1) && (iline + 1) <= 7);
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINENUMBER].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_QUANTITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_QUANTITY]) {
+ tld.generated_lineitem[LINEITEM::L_QUANTITY] = true;
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int64_t> dist(1, 50);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_QUANTITY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_quantity =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_QUANTITY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t i = 0; i < next_run; i++) {
+ // Multiply by 100 because the type is decimal(12, 2), so the decimal goes after
+ // two digits
+ int64_t quantity = dist(tld.rng) * 100;
+ l_quantity[batch_offset++] = {quantity};
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_EXTENDEDPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE]) {
+ tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+ RETURN_NOT_OK(L_QUANTITY(thread_index));
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_EXTENDEDPRICE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ const int32_t* l_partkey = reinterpret_cast<const int32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ const Decimal128* l_quantity = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->buffers[1]->data());
+ Decimal128* l_extendedprice =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ int64_t partkey = static_cast<int64_t>(l_partkey[batch_offset]);
+ // Divide by 100 to recover the integer representation (not Decimal).
+ int64_t quantity = static_cast<int64_t>(l_quantity[batch_offset]) / 100;
+
+ // Spec says to divide by 100, but that happens automatically due to this being
+ // stored to two decimal points.
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ int64_t extended_price = retail_price * quantity;
+ l_extendedprice[batch_offset] = {extended_price};
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_DISCOUNT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_DISCOUNT]) {
+ tld.generated_lineitem[LINEITEM::L_DISCOUNT] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 10);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_DISCOUNT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_discount =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_DISCOUNT]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_discount[batch_offset] = {dist(tld.rng)};
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_TAX(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_TAX]) {
+ tld.generated_lineitem[LINEITEM::L_TAX] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 8);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_TAX, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_tax = reinterpret_cast<Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_tax[batch_offset] = {dist(tld.rng)};
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RETURNFLAG(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_RETURNFLAG]) {
+ tld.generated_lineitem[LINEITEM::L_RETURNFLAG] = true;
+ RETURN_NOT_OK(L_RECEIPTDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<uint32_t> dist;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_RETURNFLAG, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_returnflag =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_receiptdate = reinterpret_cast<const uint32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ if (l_receiptdate[batch_offset] <= kCurrentDate) {
+ uint32_t r = dist(tld.rng);
+ l_returnflag[batch_offset] = (r % 2 == 1) ? 'R' : 'A';
+ } else {
+ l_returnflag[batch_offset] = 'N';
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINESTATUS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_LINESTATUS]) {
+ tld.generated_lineitem[LINEITEM::L_LINESTATUS] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_LINESTATUS, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_linestatus =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_LINESTATUS]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_shipdate =
+ reinterpret_cast<const uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ if (l_shipdate[batch_offset] > kCurrentDate)
+ l_linestatus[batch_offset] = 'O';
+ else
+ l_linestatus[batch_offset] = 'F';
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPDATE]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPDATE] = true;
+ RETURN_NOT_OK(O_ORDERDATE(thread_index));
+ const int32_t* o_orderdate = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(1, 121);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SHIPDATE, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_shipdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_shipdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMITDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_COMMITDATE]) {
+ tld.generated_lineitem[LINEITEM::L_COMMITDATE] = true;
+ const int32_t* o_orderdate = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(30, 90);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_COMMITDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_commitdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_COMMITDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_commitdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMITDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RECEIPTDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_RECEIPTDATE]) {
+ tld.generated_lineitem[LINEITEM::L_RECEIPTDATE] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(1, 30);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_RECEIPTDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_receiptdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_shipdate =
+ reinterpret_cast<const uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_receiptdate[batch_offset] = l_shipdate[batch_offset] + dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPINSTRUCT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT] = true;
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPINSTRUCT]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumInstructions - 1);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_SHIPINSTRUCT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_shipinstruct =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ const char* str = Instructions[dist(tld.rng)];
+ // Note that we don't have to memset the buffer to 0 because strncpy pads each
+ // string with 0's anyway
+ std::strncpy(l_shipinstruct + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPMODE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPMODE]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPMODE] = true;
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPMODE]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumModes - 1);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SHIPMODE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_shipmode =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_SHIPMODE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ const char* str = Modes[dist(tld.rng)];
+ std::strncpy(l_shipmode + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPMODE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_COMMENT]) {
+ tld.generated_lineitem[LINEITEM::L_COMMENT] = true;
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ // Comments are kind of sneaky: we always generate the full batch and then just
+ // bump the length
+ if (tld.lineitem[ibatch][LINEITEM::L_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.lineitem[ibatch][LINEITEM::L_COMMENT],
+ g_text.GenerateComments(batch_size_, 10, 43, tld.rng));
+ batch_offset = 0;
+ }
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ batch_offset += next_run;
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMENT].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ std::vector<Datum> orders;
+ int64_t orders_to_generate;
+ int64_t orderkey_start;
+
+ std::vector<std::vector<Datum>> lineitem;
+ std::vector<int> items_per_order;
+ int64_t lineitem_to_generate;
+ int64_t first_batch_offset;
+ std::bitset<LINEITEM::kNumCols> generated_lineitem;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex orders_output_queue_mutex_;
+ std::mutex lineitem_output_queue_mutex_;
+ std::queue<ExecBatch> orders_output_queue_;
+ std::queue<ExecBatch> lineitem_output_queue_;
+ int64_t batch_size_;
+ float scale_factor_;
+ int64_t orders_rows_to_generate_;
+ int64_t orders_rows_generated_;
+ std::vector<int> orders_cols_;
+ std::vector<int> lineitem_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<size_t> orders_batches_generated_ = {0};
+ std::atomic<size_t> lineitem_batches_generated_ = {0};
+};
+
+class SupplierGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 10000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, gen_list_));
+
+ random::pcg32_fast rng;
+ std::uniform_int_distribution<int64_t> dist(0, rows_to_generate_ - 1);
+ size_t num_special_rows = static_cast<size_t>(5 * scale_factor_);
+ std::unordered_set<int64_t> good_rows_set;
+ while (good_rows_set.size() < num_special_rows) {
+ int64_t row = dist(rng);
+ good_rows_set.insert(row);
+ }
+ std::unordered_set<int64_t> bad_rows_set;
+ while (bad_rows_set.size() < num_special_rows) {
+ int64_t bad_row;
+ do {
+ bad_row = dist(rng);
+ } while (good_rows_set.find(bad_row) != good_rows_set.end());
+ bad_rows_set.insert(bad_row);
+ }
+ good_rows_.clear();
+ bad_rows_.clear();
+ good_rows_.insert(good_rows_.end(), good_rows_set.begin(), good_rows_set.end());
+ bad_rows_.insert(bad_rows_.end(), bad_rows_set.begin(), bad_rows_set.end());
+ std::sort(good_rows_.begin(), good_rows_.end());
+ std::sort(bad_rows_.begin(), bad_rows_.end());
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) tld.rng.seed(seq);
+
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(S_SUPPKEY) \
+ F(S_NAME) \
+ F(S_ADDRESS) \
+ F(S_NATIONKEY) \
+ F(S_PHONE) \
+ F(S_ACCTBAL) \
+ F(S_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct SUPPLIER {
+ enum {
+ FOR_EACH_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) {#col, SUPPLIER::col},
+ const std::unordered_map<std::string, int> kNameMap = {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kGenerators = {FOR_EACH_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(), fixed_size_binary(25), utf8(),
+ int32(), fixed_size_binary(15), decimal(12, 2),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index) {
+ if (done_.load()) return Status::OK();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.suppkey_start = rows_generated_.fetch_add(batch_size_);
+ if (tld.suppkey_start >= rows_to_generate_) return Status::OK();
+
+ tld.to_generate = std::min(batch_size_, rows_to_generate_ - tld.suppkey_start);
+
+ tld.batch.resize(SUPPLIER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for (int col : gen_list_) RETURN_NOT_OK(kGenerators[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for (size_t i = 0; i < gen_list_.size(); i++) {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_outputted_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_outputted_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if (is_last_batch) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ return Status::OK();
+ }
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(kTypes[column], tld.to_generate, {nullptr, std::move(buff)});
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status S_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_SUPPKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_SUPPKEY));
+ int32_t* s_suppkey = reinterpret_cast<int32_t*>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ s_suppkey[irow] = static_cast<int32_t>(tld.suppkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_NAME].kind() == Datum::NONE) {
+ RETURN_NOT_OK(S_SUPPKEY(thread_index));
+ const int32_t* s_suppkey = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->data());
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NAME));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[SUPPLIER::S_NAME]);
+ char* s_name = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_NAME].array()->buffers[1]->mutable_data());
+ // Look man, I'm just following the spec ok? Section 4.2.3 as of March 1 2022
+ const char* supplier = "Supplie#r";
+ const size_t supplier_length = std::strlen(supplier);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = s_name + byte_width * irow;
+ std::strncpy(out, supplier, byte_width);
+ AppendNumberPaddedToNineDigits(out + supplier_length, s_suppkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ADDRESS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_ADDRESS].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[SUPPLIER::S_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status S_NATIONKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_NATIONKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t* s_nationkey = reinterpret_cast<int32_t*>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status S_PHONE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_PHONE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(S_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[SUPPLIER::S_PHONE]);
+ const int32_t* s_nationkey = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->data());
+ char* s_phone = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_PHONE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ GeneratePhoneNumber(s_phone + irow * byte_width, tld.rng, s_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ACCTBAL(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_ACCTBAL].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_ACCTBAL));
+ Decimal128* s_acctbal = reinterpret_cast<Decimal128*>(
+ tld.batch[SUPPLIER::S_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_acctbal[irow] = {dist(tld.rng)};
+ }
+ return Status::OK();
+ }
+
+ Status S_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[SUPPLIER::S_COMMENT],
+ g_text.GenerateComments(tld.to_generate, 25, 100, tld.rng));
+ ModifyComments(thread_index, "Recommends", good_rows_);
+ ModifyComments(thread_index, "Complaints", bad_rows_);
+ }
+ return Status::OK();
+ }
+
+ void ModifyComments(size_t thread_index, const char* review,
+ const std::vector<int64_t>& indices) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ const int32_t* offsets = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[1]->data());
+ char* str = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[2]->mutable_data());
+ const char* customer = "Customer";
+ const int32_t customer_length = static_cast<int32_t>(std::strlen(customer));
+ const int32_t review_length = static_cast<int32_t>(std::strlen(review));
+
+ auto it = std::lower_bound(indices.begin(), indices.end(), tld.suppkey_start);
+ for (; it != indices.end() && *it < tld.suppkey_start + tld.to_generate; it++) {
+ int64_t idx_in_batch = *it - tld.suppkey_start;
+ char* out = str + offsets[idx_in_batch];
+ int32_t str_length = offsets[idx_in_batch + 1] - offsets[idx_in_batch];
+ std::uniform_int_distribution<int32_t> gap_dist(
+ 0, str_length - customer_length - review_length);
+ int32_t gap = gap_dist(tld.rng);
+ int32_t total_length = customer_length + gap + review_length;
+ std::uniform_int_distribution<int32_t> start_dist(0, str_length - total_length);
+ int32_t start = start_dist(tld.rng);
+ std::memcpy(out + start, customer, customer_length);
+ std::memcpy(out + start + customer_length + gap, review, review_length);
+ }
+ }
+
+ struct ThreadLocalData {
+ random::pcg32_fast rng;
+ int64_t suppkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+ std::vector<int64_t> good_rows_;
+ std::vector<int64_t> bad_rows_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_;
+ std::atomic<int64_t> rows_generated_;
+ float scale_factor_;
+ int64_t batch_size_;
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class PartGenerator : public TpchTableGenerator {
+ public:
+ explicit PartGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetPartOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch, gen_->NextPartBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->part_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class PartSuppGenerator : public TpchTableGenerator {
+ public:
+ explicit PartSuppGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetPartSuppOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextPartSuppBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->partsupp_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class CustomerGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, gen_list_));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) tld.rng.seed(seq);
+
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(C_CUSTKEY) \
+ F(C_NAME) \
+ F(C_ADDRESS) \
+ F(C_NATIONKEY) \
+ F(C_PHONE) \
+ F(C_ACCTBAL) \
+ F(C_MKTSEGMENT) \
+ F(C_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct CUSTOMER {
+ enum {
+ FOR_EACH_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) {#col, CUSTOMER::col},
+ const std::unordered_map<std::string, int> kNameMap = {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kGenerators = {FOR_EACH_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(),
+ utf8(),
+ utf8(),
+ int32(),
+ fixed_size_binary(15),
+ decimal(12, 2),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index) {
+ if (done_.load()) return Status::OK();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.custkey_start = rows_generated_.fetch_add(batch_size_);
+ if (tld.custkey_start >= rows_to_generate_) return Status::OK();
+
+ tld.to_generate = std::min(batch_size_, rows_to_generate_ - tld.custkey_start);
+
+ tld.batch.resize(CUSTOMER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for (int col : gen_list_) RETURN_NOT_OK(kGenerators[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for (size_t i = 0; i < gen_list_.size(); i++) {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_generated_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_generated_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if (is_last_batch) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ return Status::OK();
+ }
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(kTypes[column], tld.to_generate, {nullptr, std::move(buff)});
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status C_CUSTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_CUSTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_CUSTKEY));
+ int32_t* c_custkey = reinterpret_cast<int32_t*>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ c_custkey[irow] = static_cast<int32_t>(tld.custkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_NAME].kind() == Datum::NONE) {
+ RETURN_NOT_OK(C_CUSTKEY(thread_index));
+ const int32_t* c_custkey = reinterpret_cast<const int32_t*>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->data());
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ const char* customer = "Customer#";
+ const size_t customer_length = std::strlen(customer);
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ int num_digits = GetNumDigits(c_custkey[irow]);
+ int num_chars = std::max(num_digits, 9);
+ offsets[irow + 1] =
+ static_cast<int32_t>(offsets[irow] + num_chars + customer_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[tld.to_generate]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = str + offsets[irow];
+ std::memcpy(out, customer, customer_length);
+ AppendNumberPaddedToNineDigits(out + customer_length, c_custkey[irow]);
+ }
+ ArrayData ad(utf8(), tld.to_generate,
+ {nullptr, std::move(offset_buff), std::move(str_buff)});
+ tld.batch[CUSTOMER::C_NAME] = std::move(ad);
+ }
+ return Status::OK();
+ }
+
+ Status C_ADDRESS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_ADDRESS].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[CUSTOMER::C_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status C_NATIONKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_NATIONKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t* c_nationkey = reinterpret_cast<int32_t*>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status C_PHONE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_PHONE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(C_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[CUSTOMER::C_PHONE]);
+ const int32_t* c_nationkey = reinterpret_cast<const int32_t*>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->data());
+ char* c_phone = reinterpret_cast<char*>(
+ tld.batch[CUSTOMER::C_PHONE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ GeneratePhoneNumber(c_phone + irow * byte_width, tld.rng, c_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_ACCTBAL(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_ACCTBAL].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_ACCTBAL));
+ Decimal128* c_acctbal = reinterpret_cast<Decimal128*>(
+ tld.batch[CUSTOMER::C_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_acctbal[irow] = {dist(tld.rng)};
+ }
+ return Status::OK();
+ }
+
+ Status C_MKTSEGMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_MKTSEGMENT].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_MKTSEGMENT));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[CUSTOMER::C_MKTSEGMENT]);
+ char* c_mktsegment = reinterpret_cast<char*>(
+ tld.batch[CUSTOMER::C_MKTSEGMENT].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int32_t> dist(0, kNumSegments - 1);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = c_mktsegment + irow * byte_width;
+ int str_idx = dist(tld.rng);
+ std::strncpy(out, Segments[str_idx], byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[CUSTOMER::C_COMMENT],
+ g_text.GenerateComments(tld.to_generate, 29, 116, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ random::pcg32_fast rng;
+ int64_t custkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_{0};
+ std::atomic<int64_t> rows_generated_ = {0};
+ float scale_factor_{0};
+ int64_t batch_size_{0};
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class OrdersGenerator : public TpchTableGenerator {
+ public:
+ explicit OrdersGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetOrdersOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch, gen_->NextOrdersBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->orders_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class LineitemGenerator : public TpchTableGenerator {
+ public:
+ explicit LineitemGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetLineItemOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextLineItemBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->lineitem_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class NationGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, float /*scale_factor*/,
+ int64_t /*batch_size*/) override {
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, column_indices_));
+ rng_.seed(arrow_vendored::pcg_extras::seed_seq_from<std::random_device>{});
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t /*num_threads*/, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback /*schedule_task_callback*/) override {
+ std::shared_ptr<Buffer> N_NATIONKEY_buffer =
+ Buffer::Wrap(kNationKey, sizeof(kNationKey));
+ ArrayData N_NATIONKEY_arraydata(int32(), kRowCount,
+ {nullptr, std::move(N_NATIONKEY_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> N_NAME_buffer,
+ AllocateBuffer(kRowCount * kNameByteWidth));
+ char* N_NAME = reinterpret_cast<char*>(N_NAME_buffer->mutable_data());
+ for (size_t i = 0; i < kRowCount; i++)
+ std::strncpy(N_NAME + kNameByteWidth * i, kCountryNames[i], kNameByteWidth);
+ ArrayData N_NAME_arraydata(fixed_size_binary(kNameByteWidth), kRowCount,
+ {nullptr, std::move(N_NAME_buffer)});
+
+ std::shared_ptr<Buffer> N_REGIONKEY_buffer =
+ Buffer::Wrap(kRegionKey, sizeof(kRegionKey));
+ ArrayData N_REGIONKEY_arraydata(int32(), kRowCount,
+ {nullptr, std::move(N_REGIONKEY_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(Datum N_COMMENT_datum,
+ g_text.GenerateComments(kRowCount, 31, 114, rng_));
+
+ std::vector<Datum> fields = {
+ std::move(N_NATIONKEY_arraydata), std::move(N_NAME_arraydata),
+ std::move(N_REGIONKEY_arraydata), std::move(N_COMMENT_datum)};
+
+ std::vector<Datum> result;
+ for (const int& col : column_indices_) result.push_back(fields[col]);
+ ARROW_ASSIGN_OR_RAISE(ExecBatch batch, ExecBatch::Make(std::move(result)));
+ output_callback(std::move(batch));
+ finished_callback(static_cast<int64_t>(1));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ random::pcg32_fast rng_;
+
+ static constexpr size_t kRowCount = 25;
+ static constexpr int32_t kNameByteWidth = 25;
+ const int32_t kNationKey[kRowCount] = {0, 1, 2, 3, 4, 5, 6, 7, 8,
+ 9, 10, 11, 12, 13, 14, 15, 16, 17,
+ 18, 19, 20, 21, 22, 23, 24};
+ const char* kCountryNames[kRowCount] = {
+ "ALGERIA", "ARGENTINA", "BRAZIL", "CANADA", "EGYPT",
+ "ETHIOPIA", "FRANCE", "GERMANY", "INDIA", "INDONESIA",
+ "IRAN", "IRAQ", "JAPAN", "JORDAN", "KENYA",
+ "MOROCCO", "MOZAMBIQUE", "PERU", "CHINA", "ROMANIA",
+ "SAUDI ARABIA", "VIETNAM", "RUSSIA", "UNITED KINGDOM", "UNITED STATES"};
+ const int32_t kRegionKey[kRowCount] = {0, 1, 1, 1, 4, 0, 3, 3, 2, 2, 4, 4, 2,
+ 4, 0, 0, 0, 1, 2, 3, 4, 2, 3, 3, 1};
+
+ struct NATION {
+ enum {
Review comment:
Also perhaps call this NationColumn or NationField?
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<float>(scale_factor));
+
+ ExecNode* lineitem =
+ *gen.Lineitem({"L_QUANTITY", "L_EXTENDEDPRICE", "L_TAX", "L_DISCOUNT", "L_SHIPDATE",
+ "L_RETURNFLAG", "L_LINESTATUS"});
+
+ std::shared_ptr<Date32Scalar> sept_2_1998 = std::make_shared<Date32Scalar>(
+ 10471); // September 2, 1998 is 10471 days after January 1, 1970
+ Expression filter =
+ less_equal(field_ref("L_SHIPDATE"), literal(std::move(sept_2_1998)));
+ FilterNodeOptions filter_opts(filter);
+
+ Expression l_returnflag = field_ref("L_RETURNFLAG");
+ Expression l_linestatus = field_ref("L_LINESTATUS");
+ Expression quantity = field_ref("L_QUANTITY");
+ Expression base_price = field_ref("L_EXTENDEDPRICE");
+
+ std::shared_ptr<Decimal128Scalar> decimal_1 =
+ std::make_shared<Decimal128Scalar>(Decimal128{0, 100}, decimal(12, 2));
+ Expression discount_multiplier =
+ call("subtract", {literal(decimal_1), field_ref("L_DISCOUNT")});
+ Expression tax_multiplier = call("add", {literal(decimal_1), field_ref("L_TAX")});
+ Expression disc_price =
+ call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier});
+ Expression charge =
+ call("multiply",
+ {call("cast",
+ {call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier})},
+ compute::CastOptions::Unsafe(decimal(12, 2))),
+ tax_multiplier});
+ Expression discount = field_ref("L_DISCOUNT");
+
+ std::vector<Expression> projection_list = {l_returnflag, l_linestatus, quantity,
+ base_price, disc_price, charge,
+ quantity, base_price, discount};
+ std::vector<std::string> project_names = {
+ "l_returnflag", "l_linestatus", "sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price", "avg_disc"};
+ ProjectNodeOptions project_opts(std::move(projection_list));
+
+ ScalarAggregateOptions sum_opts = ScalarAggregateOptions::Defaults();
+ CountOptions count_opts(CountOptions::CountMode::ALL);
+ std::vector<arrow::compute::internal::Aggregate> aggs = {
+ {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts},
+ {"hash_sum", &sum_opts}, {"hash_mean", &sum_opts}, {"hash_mean", &sum_opts},
+ {"hash_mean", &sum_opts}, {"hash_count", &count_opts}};
+
+ std::vector<FieldRef> cols = {2, 3, 4, 5, 6, 7, 8, 2};
Review comment:
Because one doesn't need an indirection to know what these columns are.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829337783
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
Review comment:
from google style guide:
```
Parameters are either inputs to the function, outputs from the function, or both. Non-optional input parameters should usually be values or const references, while non-optional output and input/output parameters should usually be references (which cannot be null). Generally, use std::optional to represent optional by-value inputs, and use a const pointer when the non-optional form would have used a reference. Use non-const pointers to represent optional outputs and optional input/output parameters.
```
this is a non-optional input/output parameter, so according to them it should be a reference.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829563018
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
Review comment:
As an end user of this code, I would appreciate the ability to provide determinism if that's not too painful.
We wouldn't always want or need it, but sometimes it is helpful to be able to generate the exact same set of data.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r830243120
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<float>(scale_factor));
+
+ ExecNode* lineitem =
+ *gen.Lineitem({"L_QUANTITY", "L_EXTENDEDPRICE", "L_TAX", "L_DISCOUNT", "L_SHIPDATE",
+ "L_RETURNFLAG", "L_LINESTATUS"});
+
+ std::shared_ptr<Date32Scalar> sept_2_1998 = std::make_shared<Date32Scalar>(
+ 10471); // September 2, 1998 is 10471 days after January 1, 1970
+ Expression filter =
+ less_equal(field_ref("L_SHIPDATE"), literal(std::move(sept_2_1998)));
+ FilterNodeOptions filter_opts(filter);
+
+ Expression l_returnflag = field_ref("L_RETURNFLAG");
+ Expression l_linestatus = field_ref("L_LINESTATUS");
+ Expression quantity = field_ref("L_QUANTITY");
+ Expression base_price = field_ref("L_EXTENDEDPRICE");
+
+ std::shared_ptr<Decimal128Scalar> decimal_1 =
+ std::make_shared<Decimal128Scalar>(Decimal128{0, 100}, decimal(12, 2));
+ Expression discount_multiplier =
+ call("subtract", {literal(decimal_1), field_ref("L_DISCOUNT")});
+ Expression tax_multiplier = call("add", {literal(decimal_1), field_ref("L_TAX")});
+ Expression disc_price =
+ call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier});
+ Expression charge =
+ call("multiply",
+ {call("cast",
+ {call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier})},
+ compute::CastOptions::Unsafe(decimal(12, 2))),
+ tax_multiplier});
+ Expression discount = field_ref("L_DISCOUNT");
+
+ std::vector<Expression> projection_list = {l_returnflag, l_linestatus, quantity,
+ base_price, disc_price, charge,
+ quantity, base_price, discount};
+ std::vector<std::string> project_names = {
+ "l_returnflag", "l_linestatus", "sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price", "avg_disc"};
+ ProjectNodeOptions project_opts(std::move(projection_list));
+
+ ScalarAggregateOptions sum_opts = ScalarAggregateOptions::Defaults();
+ CountOptions count_opts(CountOptions::CountMode::ALL);
+ std::vector<arrow::compute::internal::Aggregate> aggs = {
+ {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts},
+ {"hash_sum", &sum_opts}, {"hash_mean", &sum_opts}, {"hash_mean", &sum_opts},
+ {"hash_mean", &sum_opts}, {"hash_count", &count_opts}};
+
+ std::vector<FieldRef> cols = {2, 3, 4, 5, 6, 7, 8, 2};
Review comment:
ok, i don't see why that's more readable since it's mostly just a repeat of the list of names above, but i changed it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r834574097
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,631 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <regex>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+namespace internal {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+// Verifies that the data is valid Arrow and ensures it's not null.
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values) {
+ ASSERT_EQ(d.array()->buffers[0].get(), nullptr);
+ ASSERT_OK(d.make_array()->ValidateFull());
+ }
+}
+
+// Verifies that each element is seen exactly once, and that it's between min and max
+// inclusive
+void VerifyUniqueKey(std::unordered_set<int32_t>* seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen->insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const util::string_view& row, const char* prefix,
+ const int64_t i, const int32_t* nums,
+ bool verify_padding) {
+ size_t num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_TRUE(row.starts_with(prefix)) << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row.data() + num_offset;
+ int64_t num = 0;
+ size_t ibyte = num_offset;
+ // Parse the number out
+ for (; *num_str && ibyte < row.size(); ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ // If nums is not null, ensure it matches the parsed number
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ // TPC-H requires only ever requires padding up to 9 digits, so we ensure that
+ // the total length of the string was at least 9 (could be more for bigger numbers).
+ if (verify_padding) {
+ int64_t num_chars = static_cast<int64_t>(ibyte - num_offset);
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+// Verifies that each row is the string "prefix" followed by a number. If numbers is not
+// EMPTY, it also checks that the number following the prefix is equal to the
+// corresponding row in numbers. Some TPC-H data is padded to 9 zeros, which this function
+// can optionally verify as well. This string function verifies fixed width columns.
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ util::string_view view(row, byte_width);
+ VerifyStringAndNumber_Single(view, prefix, i, nums, verify_padding);
+ }
+}
+
+// Same as above but for variable length columns
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ util::string_view view(str + start, str_len);
+ VerifyStringAndNumber_Single(view, prefix, i, nums, verify_padding);
+ }
+}
+
+// Verifies that each row is a V-string, which is defined in the spec to be
+// a string of random length between min_length and max_length, that is composed
+// of alphanumeric characters, commas, or spaces.
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = off[i];
+ int32_t end = off[i + 1];
+ int32_t str_len = end - start;
+ ASSERT_LE(str_len, max_length);
+ ASSERT_GE(str_len, min_length);
+ for (int32_t i = start; i < end; i++) {
+ bool is_valid =
+ std::isdigit(str[i]) || std::isalpha(str[i]) || str[i] == ',' || str[i] == ' ';
+ ASSERT_TRUE(is_valid) << "Character " << str[i]
+ << " is not a digit, a letter, a comma, or a space";
+ }
+ }
+}
+
+// Verifies that each 32-bit element modulo "mod" is between min and max.
+void VerifyModuloBetween(const Datum& d, int32_t min, int32_t max, int32_t mod) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t m = n[i] % mod;
+ ASSERT_GE(m, min) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ ASSERT_LE(m, max) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ }
+}
+
+// Verifies that each 32-bit element is between min and max.
+void VerifyAllBetween(const Datum& d, int32_t min, int32_t max) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ ASSERT_GE(n[i], min) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ ASSERT_LE(n[i], max) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ }
+}
+
+void VerifyNationKey(const Datum& d) { VerifyAllBetween(d, 0, 24); }
+
+// Verifies that each row satisfies the phone number spec.
+void VerifyPhone(const Datum& d) {
+ int64_t length = d.length();
+ const char* phones = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ constexpr int kByteWidth = 15; // This is common for all PHONE columns
+ std::regex exp("\\d{2}-\\d{3}-\\d{3}-\\d{4}");
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = phones + i * kByteWidth;
+ ASSERT_TRUE(std::regex_match(row, row + kByteWidth, exp));
+ }
+}
+
+// Verifies that each decimal is between min and max
+void VerifyDecimalsBetween(const Datum& d, int64_t min, int64_t max) {
+ int64_t length = d.length();
+ const Decimal128* decs =
+ reinterpret_cast<const Decimal128*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int64_t val = static_cast<int64_t>(decs[i]);
+ ASSERT_LE(val, max);
+ ASSERT_GE(val, min);
+ }
+}
+
+// Verifies that each variable-length row is a series of words separated by
+// spaces. Number of words is determined by the number of spaces.
+void VerifyCorrectNumberOfWords_Varlen(const Datum& d, int num_words) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+
+ int32_t start = offsets[i];
+ int32_t end = offsets[i + 1];
+ int32_t str_len = end - start;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, str + start, str_len);
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = offsets[i]; j < offsets[i + 1]; j++) {
+ bool is_space = str[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(str[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << tmp_str;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << tmp_str;
+ }
+}
+
+// Same as above but for fixed width columns.
+void VerifyCorrectNumberOfWords_FixedWidth(const Datum& d, int num_words,
+ int byte_width) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+ const char* row = str + i * byte_width;
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = 0; j < byte_width && row[j]; j++) {
+ bool is_space = row[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(row[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << row;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << row;
+ }
+}
+
+// Verifies that each row of the single-byte-wide column is one of the possibilities.
+void VerifyOneOf(const Datum& d, const std::unordered_set<char>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++)
+ ASSERT_TRUE(possibilities.find(col[i]) != possibilities.end());
+}
+
+// Verifies that each fixed-width row is one of the possibilities
+void VerifyOneOf(const Datum& d, int32_t byte_width,
+ const std::unordered_set<util::string_view>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = col + i * byte_width;
+ int32_t row_len = 0;
+ while (row[row_len] && row_len < byte_width) row_len++;
+ util::string_view view(row, row_len);
+ ASSERT_TRUE(possibilities.find(view) != possibilities.end())
+ << view << " is not a valid string.";
+ }
+}
+
+// Counts the number of instances of each integer
+void CountInstances(std::unordered_map<int32_t, int32_t>& counts, const Datum& d) {
+ int64_t length = d.length();
+ const int32_t* nums = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) counts[nums[i]]++;
+}
+
+// For the S_COMMENT column, some of the columns must be modified to contain
+// "Customer...Complaints" or "Customer...Recommends". This function counts the number of
+// good and bad comments.
+void CountModifiedComments(const Datum& d, int& good_count, int& bad_count) {
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ // Length of S_COMMENT is at most 100
+ char tmp_string[101];
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + offsets[i];
+ int32_t row_length = offsets[i + 1] - offsets[i];
+ std::memset(tmp_string, 0, sizeof(tmp_string));
+ std::memcpy(tmp_string, row, row_length);
+ char* customer = std::strstr(tmp_string, "Customer");
+ char* recommends = std::strstr(tmp_string, "Recommends");
+ char* complaints = std::strstr(tmp_string, "Complaints");
+ if (customer) {
+ ASSERT_TRUE((recommends != nullptr) ^ (complaints != nullptr));
+ if (recommends) good_count++;
+ if (complaints) bad_count++;
+ }
+ }
+}
+
+TEST(TpchNode, ScaleFactor) {
+ ExecContext ctx(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(&ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), 0.25f);
Review comment:
test fixture is usually used for if you need the same object to live between multiple tests, which we don't, but i agree there is a bit of repeated code. i've factored it out into a function.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r834553191
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,631 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <regex>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+namespace internal {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+// Verifies that the data is valid Arrow and ensures it's not null.
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values) {
+ ASSERT_EQ(d.array()->buffers[0].get(), nullptr);
+ ASSERT_OK(d.make_array()->ValidateFull());
+ }
+}
+
+// Verifies that each element is seen exactly once, and that it's between min and max
+// inclusive
+void VerifyUniqueKey(std::unordered_set<int32_t>* seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen->insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const util::string_view& row, const char* prefix,
+ const int64_t i, const int32_t* nums,
+ bool verify_padding) {
+ size_t num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_TRUE(row.starts_with(prefix)) << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row.data() + num_offset;
+ int64_t num = 0;
+ size_t ibyte = num_offset;
+ // Parse the number out
+ for (; *num_str && ibyte < row.size(); ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ // If nums is not null, ensure it matches the parsed number
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ // TPC-H requires only ever requires padding up to 9 digits, so we ensure that
+ // the total length of the string was at least 9 (could be more for bigger numbers).
+ if (verify_padding) {
+ int64_t num_chars = static_cast<int64_t>(ibyte - num_offset);
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+// Verifies that each row is the string "prefix" followed by a number. If numbers is not
+// EMPTY, it also checks that the number following the prefix is equal to the
+// corresponding row in numbers. Some TPC-H data is padded to 9 zeros, which this function
+// can optionally verify as well. This string function verifies fixed width columns.
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ util::string_view view(row, byte_width);
+ VerifyStringAndNumber_Single(view, prefix, i, nums, verify_padding);
+ }
+}
+
+// Same as above but for variable length columns
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ util::string_view view(str + start, str_len);
+ VerifyStringAndNumber_Single(view, prefix, i, nums, verify_padding);
+ }
+}
+
+// Verifies that each row is a V-string, which is defined in the spec to be
+// a string of random length between min_length and max_length, that is composed
+// of alphanumeric characters, commas, or spaces.
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = off[i];
+ int32_t end = off[i + 1];
+ int32_t str_len = end - start;
+ ASSERT_LE(str_len, max_length);
+ ASSERT_GE(str_len, min_length);
+ for (int32_t i = start; i < end; i++) {
+ bool is_valid =
+ std::isdigit(str[i]) || std::isalpha(str[i]) || str[i] == ',' || str[i] == ' ';
+ ASSERT_TRUE(is_valid) << "Character " << str[i]
+ << " is not a digit, a letter, a comma, or a space";
+ }
+ }
+}
+
+// Verifies that each 32-bit element modulo "mod" is between min and max.
+void VerifyModuloBetween(const Datum& d, int32_t min, int32_t max, int32_t mod) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t m = n[i] % mod;
+ ASSERT_GE(m, min) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ ASSERT_LE(m, max) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ }
+}
+
+// Verifies that each 32-bit element is between min and max.
+void VerifyAllBetween(const Datum& d, int32_t min, int32_t max) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ ASSERT_GE(n[i], min) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ ASSERT_LE(n[i], max) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ }
+}
+
+void VerifyNationKey(const Datum& d) { VerifyAllBetween(d, 0, 24); }
+
+// Verifies that each row satisfies the phone number spec.
+void VerifyPhone(const Datum& d) {
+ int64_t length = d.length();
+ const char* phones = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ constexpr int kByteWidth = 15; // This is common for all PHONE columns
+ std::regex exp("\\d{2}-\\d{3}-\\d{3}-\\d{4}");
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = phones + i * kByteWidth;
+ ASSERT_TRUE(std::regex_match(row, row + kByteWidth, exp));
+ }
+}
+
+// Verifies that each decimal is between min and max
+void VerifyDecimalsBetween(const Datum& d, int64_t min, int64_t max) {
+ int64_t length = d.length();
+ const Decimal128* decs =
+ reinterpret_cast<const Decimal128*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int64_t val = static_cast<int64_t>(decs[i]);
+ ASSERT_LE(val, max);
+ ASSERT_GE(val, min);
+ }
+}
+
+// Verifies that each variable-length row is a series of words separated by
+// spaces. Number of words is determined by the number of spaces.
+void VerifyCorrectNumberOfWords_Varlen(const Datum& d, int num_words) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+
+ int32_t start = offsets[i];
+ int32_t end = offsets[i + 1];
+ int32_t str_len = end - start;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, str + start, str_len);
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = offsets[i]; j < offsets[i + 1]; j++) {
+ bool is_space = str[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(str[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << tmp_str;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << tmp_str;
+ }
+}
+
+// Same as above but for fixed width columns.
+void VerifyCorrectNumberOfWords_FixedWidth(const Datum& d, int num_words,
+ int byte_width) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+ const char* row = str + i * byte_width;
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = 0; j < byte_width && row[j]; j++) {
+ bool is_space = row[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(row[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << row;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << row;
+ }
+}
+
+// Verifies that each row of the single-byte-wide column is one of the possibilities.
+void VerifyOneOf(const Datum& d, const std::unordered_set<char>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++)
+ ASSERT_TRUE(possibilities.find(col[i]) != possibilities.end());
+}
+
+// Verifies that each fixed-width row is one of the possibilities
+void VerifyOneOf(const Datum& d, int32_t byte_width,
+ const std::unordered_set<util::string_view>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = col + i * byte_width;
+ int32_t row_len = 0;
+ while (row[row_len] && row_len < byte_width) row_len++;
+ util::string_view view(row, row_len);
+ ASSERT_TRUE(possibilities.find(view) != possibilities.end())
+ << view << " is not a valid string.";
+ }
+}
+
+// Counts the number of instances of each integer
+void CountInstances(std::unordered_map<int32_t, int32_t>& counts, const Datum& d) {
Review comment:
Ok, i think i got all of them
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r836846192
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,123 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <benchmark/benchmark.h>
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+namespace internal {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<double>(scale_factor));
+
+ ExecNode* lineitem =
+ *gen.Lineitem({"L_QUANTITY", "L_EXTENDEDPRICE", "L_TAX", "L_DISCOUNT", "L_SHIPDATE",
+ "L_RETURNFLAG", "L_LINESTATUS"});
+
+ auto sept_2_1998 = std::make_shared<Date32Scalar>(
+ 10471); // September 2, 1998 is 10471 days after January 1, 1970
+ Expression filter =
+ less_equal(field_ref("L_SHIPDATE"), literal(std::move(sept_2_1998)));
+ FilterNodeOptions filter_opts(filter);
+
+ Expression l_returnflag = field_ref("L_RETURNFLAG");
+ Expression l_linestatus = field_ref("L_LINESTATUS");
+ Expression quantity = field_ref("L_QUANTITY");
+ Expression base_price = field_ref("L_EXTENDEDPRICE");
+
+ std::shared_ptr<Decimal128Scalar> decimal_1 =
+ std::make_shared<Decimal128Scalar>(Decimal128{0, 100}, decimal(12, 2));
+ Expression discount_multiplier =
+ call("subtract", {literal(decimal_1), field_ref("L_DISCOUNT")});
+ Expression tax_multiplier = call("add", {literal(decimal_1), field_ref("L_TAX")});
+ Expression disc_price =
+ call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier});
+ Expression charge =
+ call("multiply",
+ {call("cast",
+ {call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier})},
+ compute::CastOptions::Unsafe(decimal(12, 2))),
+ tax_multiplier});
+ Expression discount = field_ref("L_DISCOUNT");
+
+ std::vector<Expression> projection_list = {l_returnflag, l_linestatus, quantity,
+ base_price, disc_price, charge,
+ quantity, base_price, discount};
+ std::vector<std::string> project_names = {
+ "l_returnflag", "l_linestatus", "sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price", "avg_disc"};
+ ProjectNodeOptions project_opts(std::move(projection_list), std::move(project_names));
+
+ ScalarAggregateOptions sum_opts = ScalarAggregateOptions::Defaults();
+ CountOptions count_opts(CountOptions::CountMode::ALL);
+ std::vector<arrow::compute::internal::Aggregate> aggs = {
+ {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts},
+ {"hash_sum", &sum_opts}, {"hash_mean", &sum_opts}, {"hash_mean", &sum_opts},
+ {"hash_mean", &sum_opts}, {"hash_count", &count_opts}};
+
+ std::vector<FieldRef> to_aggregate = {"sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price",
+ "avg_disc", "sum_qty"};
+
+ std::vector<std::string> names = {"sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price",
+ "avg_disc", "count_order"};
+
+ std::vector<FieldRef> keys = {"l_returnflag", "l_linestatus"};
+ AggregateNodeOptions agg_opts(aggs, to_aggregate, names, keys);
+
+ SortKey l_returnflag_key("l_returnflag");
+ SortKey l_linestatus_key("l_linestatus");
+ SortOptions sort_opts({l_returnflag_key, l_linestatus_key});
+ OrderBySinkNodeOptions order_by_opts(sort_opts, sink_gen);
+
+ Declaration filter_decl("filter", {Declaration::Input(lineitem)}, filter_opts);
+ Declaration project_decl("project", project_opts);
+ Declaration aggregate_decl("aggregate", agg_opts);
+ Declaration orderby_decl("order_by_sink", order_by_opts);
+
+ Declaration q1 =
+ Declaration::Sequence({filter_decl, project_decl, aggregate_decl, orderby_decl});
+ std::ignore = *q1.AddToPlan(plan.get());
+ return plan;
+}
+
+static void BM_Tpch_Q1(benchmark::State& st) {
+ for (auto _ : st) {
+ st.PauseTiming();
+ AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+ std::shared_ptr<ExecPlan> plan = Plan_Q1(&sink_gen, static_cast<int>(st.range(0)));
+ st.ResumeTiming();
+ auto fut = StartAndCollect(plan.get(), sink_gen);
+ auto res = *fut.MoveResult();
+ }
+}
Review comment:
I'm not sure the benchmarks will be entirely deleted. The separate tool is for profiling and not benchmarking. Ideally we can have one "builtin queries" spot where these queries are defined and then we can use them in both the query tester (for profiling) and the benchmarks (e.g. for detecting regressions).
I agree that output items per second is probably meaningless. In the long term I think we would want something like "(input rows * cost_factor) per second" where `cost_factor` is some kind of empirically calculated estimate based on how long a query "should" take based on the complexity. E.g. a scan-filter-project-sink would have a relatively low cost factor but a query with joins would have a higher cost factor.
That's somewhat theoretical though. Until then I think a relatively intuitive and meaningful metric would be "input bytes per second". It may not be a simple thing to add though since the size of the input is somewhat random and so we'd need to keep some kind of counters while running.
Let's add a follow-up PR to address this more thoroughly. It can just be "add TPC-H benchmarks to the benchmarks suite".
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r837887327
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,123 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <benchmark/benchmark.h>
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+namespace internal {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<double>(scale_factor));
+
+ ExecNode* lineitem =
+ *gen.Lineitem({"L_QUANTITY", "L_EXTENDEDPRICE", "L_TAX", "L_DISCOUNT", "L_SHIPDATE",
+ "L_RETURNFLAG", "L_LINESTATUS"});
+
+ auto sept_2_1998 = std::make_shared<Date32Scalar>(
+ 10471); // September 2, 1998 is 10471 days after January 1, 1970
+ Expression filter =
+ less_equal(field_ref("L_SHIPDATE"), literal(std::move(sept_2_1998)));
+ FilterNodeOptions filter_opts(filter);
+
+ Expression l_returnflag = field_ref("L_RETURNFLAG");
+ Expression l_linestatus = field_ref("L_LINESTATUS");
+ Expression quantity = field_ref("L_QUANTITY");
+ Expression base_price = field_ref("L_EXTENDEDPRICE");
+
+ std::shared_ptr<Decimal128Scalar> decimal_1 =
+ std::make_shared<Decimal128Scalar>(Decimal128{0, 100}, decimal(12, 2));
+ Expression discount_multiplier =
+ call("subtract", {literal(decimal_1), field_ref("L_DISCOUNT")});
+ Expression tax_multiplier = call("add", {literal(decimal_1), field_ref("L_TAX")});
+ Expression disc_price =
+ call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier});
+ Expression charge =
+ call("multiply",
+ {call("cast",
+ {call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier})},
+ compute::CastOptions::Unsafe(decimal(12, 2))),
+ tax_multiplier});
+ Expression discount = field_ref("L_DISCOUNT");
+
+ std::vector<Expression> projection_list = {l_returnflag, l_linestatus, quantity,
+ base_price, disc_price, charge,
+ quantity, base_price, discount};
+ std::vector<std::string> project_names = {
+ "l_returnflag", "l_linestatus", "sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price", "avg_disc"};
+ ProjectNodeOptions project_opts(std::move(projection_list), std::move(project_names));
+
+ ScalarAggregateOptions sum_opts = ScalarAggregateOptions::Defaults();
+ CountOptions count_opts(CountOptions::CountMode::ALL);
+ std::vector<arrow::compute::internal::Aggregate> aggs = {
+ {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts},
+ {"hash_sum", &sum_opts}, {"hash_mean", &sum_opts}, {"hash_mean", &sum_opts},
+ {"hash_mean", &sum_opts}, {"hash_count", &count_opts}};
+
+ std::vector<FieldRef> to_aggregate = {"sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price",
+ "avg_disc", "sum_qty"};
+
+ std::vector<std::string> names = {"sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price",
+ "avg_disc", "count_order"};
+
+ std::vector<FieldRef> keys = {"l_returnflag", "l_linestatus"};
+ AggregateNodeOptions agg_opts(aggs, to_aggregate, names, keys);
+
+ SortKey l_returnflag_key("l_returnflag");
+ SortKey l_linestatus_key("l_linestatus");
+ SortOptions sort_opts({l_returnflag_key, l_linestatus_key});
+ OrderBySinkNodeOptions order_by_opts(sort_opts, sink_gen);
+
+ Declaration filter_decl("filter", {Declaration::Input(lineitem)}, filter_opts);
+ Declaration project_decl("project", project_opts);
+ Declaration aggregate_decl("aggregate", agg_opts);
+ Declaration orderby_decl("order_by_sink", order_by_opts);
+
+ Declaration q1 =
+ Declaration::Sequence({filter_decl, project_decl, aggregate_decl, orderby_decl});
+ std::ignore = *q1.AddToPlan(plan.get());
+ return plan;
+}
+
+static void BM_Tpch_Q1(benchmark::State& st) {
+ for (auto _ : st) {
+ st.PauseTiming();
+ AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+ std::shared_ptr<ExecPlan> plan = Plan_Q1(&sink_gen, static_cast<int>(st.range(0)));
+ st.ResumeTiming();
+ auto fut = StartAndCollect(plan.get(), sink_gen);
+ auto res = *fut.MoveResult();
+ }
+}
Review comment:
It makes sense in the same way the total time makes sense but it is hard to draw any conclusions from such a number. For example, we can't easily compare it with our parsing & decoding times to know if our query is an order of magnitude slower or faster than, for example, reading a CSV file or reading a parquet file.
If we want, for example, 1GiBps throughout through some system then it would be nice to get a general sense for which parts of the system will be too slow or the most likely bottlenecks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] ursabot edited a comment on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
ursabot edited a comment on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082832095
Benchmark runs are scheduled for baseline = 0deefd1838bf4706232b90b72da010698010f7b6 and contender = 50fab73da58c9f76fbe675d661e56087e2f5fb52. 50fab73da58c9f76fbe675d661e56087e2f5fb52 is a master commit associated with this PR. Results will be available as each benchmark for each run completes.
Conbench compare runs links:
[Finished :arrow_down:0.0% :arrow_up:0.0%] [ec2-t3-xlarge-us-east-2](https://conbench.ursa.dev/compare/runs/3f8ac59fddef426cac31b6b2826ce5a6...e5b78c70b6fe4a699c4cb8b967fe5736/)
[Finished :arrow_down:0.38% :arrow_up:0.33%] [test-mac-arm](https://conbench.ursa.dev/compare/runs/02b74ab900d440f8b5361300f15ad15c...b870e78b617e41e98c2ade7241e0aa06/)
[Scheduled] [ursa-i9-9960x](https://conbench.ursa.dev/compare/runs/1f7740546747452aa61fece1295aa0e3...f66d8edd43724d9988628cbe4a4e4d50/)
[Finished :arrow_down:0.04% :arrow_up:0.26%] [ursa-thinkcentre-m75q](https://conbench.ursa.dev/compare/runs/c3c55d6999fe4d7eabec73cbdfa1fad6...4b9f45f6442447d3a9d758598d0e14d3/)
Buildkite builds:
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/415| `50fab73d` ec2-t3-xlarge-us-east-2>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-test-mac-arm/builds/401| `50fab73d` test-mac-arm>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-thinkcentre-m75q/builds/411| `50fab73d` ursa-thinkcentre-m75q>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/414| `0deefd18` ec2-t3-xlarge-us-east-2>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-test-mac-arm/builds/400| `0deefd18` test-mac-arm>
[Scheduled] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-i9-9960x/builds/400| `0deefd18` ursa-i9-9960x>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-thinkcentre-m75q/builds/410| `0deefd18` ursa-thinkcentre-m75q>
Supported benchmarks:
ec2-t3-xlarge-us-east-2: Supported benchmark langs: Python, R. Runs only benchmarks with cloud = True
test-mac-arm: Supported benchmark langs: C++, Python, R
ursa-i9-9960x: Supported benchmark langs: Python, R, JavaScript
ursa-thinkcentre-m75q: Supported benchmark langs: C++, Java
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r836848688
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3483 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/formatting.h"
+#include "arrow/util/future.h"
+#include "arrow/util/io_util.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+using internal::GetRandomSeed;
+
+namespace compute {
+namespace internal {
+/*
+Architecture of the generator:
+ This is a multithreaded implementation of TPC-H's DBGen data generator. For each table
+ that doesn't depend on any other tables, it gets its own generator. For tables that
+ are dependent on each other (namely [ORDERS, LINEITEM] and [PART, PARTSUPP]), we
+implement a "double generator". Each generator is given a list of columns to generate.
+Columns are then generated lazily a batch at a time. If column A depends on another column
+B, column A will ensure that column B is always generated. We don't have to worry about
+breaking cycles of dependencies because the dependency graph is acyclic (so there's no
+code to even bother breaking cycles). Double generators work by maintaining two output
+queues. Batches for each of the tables are generated in sync, and batches that are not
+being output immediately are appended to the queue.
+
+ To generate a batch, we grab a lock, increment a counter, check the output queue (if
+applicable), and then call the necessary generator functions. To generate a column, we
+first check if it's already been generated. If not, we allocate the batch and then fill it
+according to the spec.
+
+ There are a few types of columns that get generated:
+ - Primary Keys: incrementing counters from 1 to N (with N being the number of rows).
+These are generated by incrementing a counter (this counter is gated under a lock).
+ - V-String: Random-length string of alphanumerics
+ - Phone Number: 2, 3, 3, and 4, digit numbers separated by -'s
+ - Random Numbers: random numbers within some range
+ - Expressions: expressions based on some other columns
+ Please consult the spec for intended behavior of each individual column. Columns are
+generated by a function of the same name (so e.g. the column PS_PARTKEY is generated by a
+function PS_PARTKEY).
+*/
+
+namespace {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+std::uniform_int_distribution<int64_t> kSeedDist(std::numeric_limits<int64_t>::min(),
+ std::numeric_limits<int64_t>::max());
+} // namespace
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, double scale_factor,
+ int64_t batch_size, int64_t seed) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ int64_t seed_ = {0};
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ Unreachable("Positive 64-bit integer should never have more than 18 digits");
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ size_t kPad = 9;
+ out += std::max(kPad, static_cast<size_t>(GetNumDigits(x)));
+ arrow::internal::detail::FormatAllDigitsLeftPadded(x, kPad, '0', &out);
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+
+ out += 15;
+ arrow::internal::detail::FormatAllDigits(l3, &out);
+ *(--out) = '-';
+ arrow::internal::detail::FormatAllDigits(l2, &out);
+ *(--out) = '-';
+ arrow::internal::detail::FormatAllDigits(l1, &out);
+ *(--out) = '-';
+ arrow::internal::detail::FormatTwoDigits(country_code, &out);
+}
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, double scale_factor, int64_t seed) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ thread_local_data_.resize(num_threads);
+ random::pcg64_fast seed_rng(seed);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(kSeedDist(seed_rng));
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ if (!partsupp_output_queue_.empty()) {
+ ExecBatch result = std::move(partsupp_output_queue_.front());
+ partsupp_output_queue_.pop();
+ return std::move(result);
+ }
+ }
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+ if (!part_cols_.empty()) {
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch part_batch,
+ ExecBatch::Make(std::move(part_result)));
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ part_output_queue_.emplace(std::move(part_batch));
+ }
+ }
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ for (size_t i = 1; i < partsupp_results.size(); i++)
+ partsupp_output_queue_.emplace(std::move(partsupp_results[i]));
+ }
+ return std::move(partsupp_results[0]);
+ }
+
+ private:
+#define FOR_EACH_PART_COLUMN(F) \
+ F(P_PARTKEY) \
+ F(P_NAME) \
+ F(P_MFGR) \
+ F(P_BRAND) \
+ F(P_TYPE) \
+ F(P_SIZE) \
+ F(P_CONTAINER) \
+ F(P_RETAILPRICE) \
+ F(P_COMMENT)
+
+#define FOR_EACH_PARTSUPP_COLUMN(F) \
+ F(PS_PARTKEY) \
+ F(PS_SUPPKEY) \
+ F(PS_AVAILQTY) \
+ F(PS_SUPPLYCOST) \
+ F(PS_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct PART {
+ enum {
+ FOR_EACH_PART_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct PARTSUPP {
+ enum {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, PART::col},
+ const std::unordered_map<std::string, int> kPartNameMap = {
+ FOR_EACH_PART_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, PARTSUPP::col},
+ const std::unordered_map<std::string, int> kPartsuppNameMap = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kPartGenerators = {FOR_EACH_PART_COLUMN(MAKE_FN_ARRAY)};
+ std::vector<GenerateColumnFn> kPartsuppGenerators = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kPartTypes = {
+ int32(), utf8(), fixed_size_binary(25), fixed_size_binary(10),
+ utf8(), int32(), fixed_size_binary(10), decimal(12, 2),
+ utf8(),
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kPartsuppTypes = {
+ int32(), int32(), int32(), decimal(12, 2), utf8(),
+ };
+
+ Status AllocatePartBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.part[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.part_to_generate * byte_width));
+ ArrayData ad(kPartTypes[column], tld.part_to_generate, {nullptr, std::move(buff)});
+ tld.part[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status P_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_PARTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_PARTKEY));
+ int32_t* p_partkey = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) {
+ p_partkey[i] = static_cast<int32_t>(tld.partkey_start + i + 1);
+ ARROW_DCHECK(1 <= p_partkey[i] && p_partkey[i] <= part_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_NAME].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(
+ 0, static_cast<uint8_t>(kNumNameParts - 1));
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dist(tld.rng));
+ tld.string_indices[irow * 5 + ipart] = name_part_index;
+ string_length += std::strlen(NameParts[name_part_index]);
+ }
+ // Add 4 because there is a space between each word (i.e. four spaces)
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length + 4);
+ }
+ // Add an extra byte for the space after in the very last string.
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate] + 1));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 5 + ipart];
+ const char* part = NameParts[name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ *row++ = ' ';
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_NAME], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_NAME] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_MFGR(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_MFGR].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* manufacturer = "Manufacturer#";
+ const size_t manufacturer_length = std::strlen(manufacturer);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_MFGR));
+ char* p_mfgr = reinterpret_cast<char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ std::strncpy(p_mfgr + byte_width * irow, manufacturer, byte_width);
+ char mfgr_id = '0' + dist(tld.rng);
+ *(p_mfgr + byte_width * irow + manufacturer_length) = mfgr_id;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_BRAND(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_BRAND].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_MFGR(thread_index));
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* brand = "Brand#";
+ const size_t brand_length = std::strlen(brand);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_BRAND));
+ const char* p_mfgr = reinterpret_cast<const char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->data());
+ char* p_brand = reinterpret_cast<char*>(
+ tld.part[PART::P_BRAND].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_BRAND]);
+ int32_t mfgr_byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ const size_t mfgr_id_offset = std::strlen("Manufacturer#");
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = p_brand + byte_width * irow;
+ char mfgr_id = *(p_mfgr + irow * mfgr_byte_width + mfgr_id_offset);
+ char brand_id = '0' + dist(tld.rng);
+ std::strncpy(row, brand, byte_width);
+ *(row + brand_length) = mfgr_id;
+ *(row + brand_length + 1) = brand_id;
+ irow += 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_TYPE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_TYPE].kind() == Datum::NONE) {
+ using D = std::uniform_int_distribution<uint32_t>;
+ D dists[] = {
+ D{0, static_cast<uint8_t>(kNumTypes_1 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_2 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_3 - 1)},
+ };
+
+ const char** types[] = {Types_1, Types_2, Types_3};
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dists[ipart](tld.rng));
+ tld.string_indices[irow * 3 + ipart] = name_part_index;
+ string_length += std::strlen(types[ipart][name_part_index]);
+ }
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate]));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 3 + ipart];
+ const char* part = types[ipart][name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_TYPE], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_TYPE] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_SIZE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_SIZE].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(1, 50);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_SIZE));
+ int32_t* p_size = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_SIZE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) p_size[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status P_CONTAINER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_CONTAINER].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist1(
+ 0, static_cast<uint8_t>(kNumContainers_1 - 1));
+ std::uniform_int_distribution<int> dist2(
+ 0, static_cast<uint8_t>(kNumContainers_2 - 1));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_CONTAINER));
+ char* p_container = reinterpret_cast<char*>(
+ tld.part[PART::P_CONTAINER].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_CONTAINER]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int container1_idx = dist1(tld.rng);
+ int container2_idx = dist2(tld.rng);
+ const char* container1 = Containers_1[container1_idx];
+ const char* container2 = Containers_2[container2_idx];
+ size_t container1_length = std::strlen(container1);
+ size_t container2_length = std::strlen(container2);
+
+ char* row = p_container + byte_width * irow;
+ std::strncpy(row, container1, byte_width);
+ std::memcpy(row + container1_length, container2, container2_length);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_RETAILPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_RETAILPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_RETAILPRICE));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+ Decimal128* p_retailprice = reinterpret_cast<Decimal128*>(
+ tld.part[PART::P_RETAILPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int32_t partkey = p_partkey[irow];
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ p_retailprice[irow] = {retail_price};
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.part[PART::P_COMMENT],
+ g_text.GenerateComments(tld.part_to_generate, 5, 22, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ int64_t PartsuppBatchesToGenerate(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t num_batches = (ps_to_generate + batch_size_ - 1) / batch_size_;
+ return num_batches;
+ }
+
+ Status InitPartsupp(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.generated_partsupp.reset();
+ int64_t num_batches = PartsuppBatchesToGenerate(thread_index);
+ tld.partsupp.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.partsupp) {
+ batch.resize(PARTSUPP::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocatePartSuppBatch(size_t thread_index, size_t ibatch, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartsuppTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kPartsuppTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.partsupp[ibatch][column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status PS_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_PARTKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_PARTKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_PARTKEY));
+ int32_t* ps_partkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run; ipartsupp++, irun++)
+ ps_partkey[batch_offset++] = p_partkey[ipart];
+
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPKEY));
+ int32_t* ps_suppkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run;
+ ipartsupp++, irun++) {
+ int32_t supplier = static_cast<int32_t>(ipartsupp);
+ int32_t partkey = p_partkey[ipart];
+ // Magic formula from TPCH spec.
+ ps_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_AVAILQTY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_AVAILQTY]) {
+ tld.generated_partsupp[PARTSUPP::PS_AVAILQTY] = true;
+ std::uniform_int_distribution<int32_t> dist(1, 9999);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_AVAILQTY));
+ int32_t* ps_availqty =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for (int64_t irun = 0; irun < next_run; irun++) ps_availqty[irun] = dist(tld.rng);
+
+ tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPLYCOST(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST] = true;
+ std::uniform_int_distribution<int64_t> dist(100, 100000);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(
+ AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPLYCOST));
+ Decimal128* ps_supplycost =
+ reinterpret_cast<Decimal128*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for (int64_t irun = 0; irun < next_run; irun++)
+ ps_supplycost[irun] = {dist(tld.rng)};
+
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PARTSUPP::PS_COMMENT].kind() == Datum::NONE) {
+ int64_t irow = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ int64_t num_rows = std::min(batch_size_, ps_to_generate - irow);
+ ARROW_ASSIGN_OR_RAISE(tld.partsupp[ibatch][PARTSUPP::PS_COMMENT],
+ g_text.GenerateComments(num_rows, 49, 198, tld.rng));
+ irow += num_rows;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ std::vector<Datum> part;
+ std::vector<int8_t> string_indices;
+ int64_t part_to_generate{0};
+ int64_t partkey_start{0};
+
+ std::vector<std::vector<Datum>> partsupp;
+ std::bitset<PARTSUPP::kNumCols> generated_partsupp;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex part_output_queue_mutex_;
+ std::mutex partsupp_output_queue_mutex_;
+ std::queue<ExecBatch> part_output_queue_;
+ std::queue<ExecBatch> partsupp_output_queue_;
+ int64_t batch_size_{0};
+ double scale_factor_{0};
+ int64_t part_rows_to_generate_{0};
+ int64_t part_rows_generated_{0};
+ std::vector<int> part_cols_;
+ std::vector<int> partsupp_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<int64_t> part_batches_generated_ = {0};
+ std::atomic<int64_t> partsupp_batches_generated_ = {0};
+ static constexpr int64_t kPartSuppRowsPerPart = 4;
+};
+
+class OrdersAndLineItemGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, double scale_factor, int64_t seed) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ thread_local_data_.resize(num_threads);
+ random::pcg64_fast seed_rng(seed);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ tld.items_per_order.resize(batch_size_);
+ tld.rng.seed(kSeedDist(seed_rng));
+ }
+ orders_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000 * 10);
+ }
+ return Status::OK();
+ }
+
+ int64_t orders_batches_generated() const { return orders_batches_generated_.load(); }
+
+ int64_t lineitem_batches_generated() const {
+ return lineitem_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetOrdersOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kOrdersTypes, kOrdersNameMap, orders_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetLineItemOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kLineitemTypes, kLineitemNameMap, lineitem_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextOrdersBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if (!orders_output_queue_.empty()) {
+ ExecBatch batch = std::move(orders_output_queue_.front());
+ orders_output_queue_.pop();
+ return std::move(batch);
+ } else if (orders_rows_generated_ == orders_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate =
+ std::min(batch_size_, orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.first_batch_offset = 0;
+ tld.generated_lineitem.reset();
+
+ for (int col : orders_cols_) RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for (int col : lineitem_cols_) RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for (size_t i = 0; i < orders_cols_.size(); i++) {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ if (!lineitem_cols_.empty()) {
+ std::vector<ExecBatch> lineitem_results;
+ for (size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++) {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for (size_t icol = 0; icol < lineitem_cols_.size(); icol++) {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(lineitem_output_queue_mutex_);
+ for (ExecBatch& eb : lineitem_results) {
+ lineitem_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(orders_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextLineItemBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ExecBatch queued;
+ bool from_queue = false;
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ if (!lineitem_output_queue_.empty()) {
+ queued = std::move(lineitem_output_queue_.front());
+ lineitem_output_queue_.pop();
+ from_queue = true;
+ }
+ }
+ tld.first_batch_offset = 0;
+ if (from_queue) {
+ ARROW_DCHECK(queued.length <= batch_size_);
+ tld.first_batch_offset = queued.length;
+ if (queued.length == batch_size_) return std::move(queued);
+ }
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if (orders_rows_generated_ == orders_rows_to_generate_) {
+ if (from_queue) return std::move(queued);
+ return util::nullopt;
+ }
+
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate =
+ std::min(batch_size_, orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1ll);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.generated_lineitem.reset();
+ if (from_queue) {
+ lineitem_batches_generated_.fetch_sub(1);
+ for (size_t i = 0; i < lineitem_cols_.size(); i++)
+ if (tld.lineitem[0][lineitem_cols_[i]].kind() == Datum::NONE)
+ tld.lineitem[0][lineitem_cols_[i]] = std::move(queued[i]);
+ }
+
+ for (int col : orders_cols_) RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for (int col : lineitem_cols_) RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ if (!orders_cols_.empty()) {
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for (size_t i = 0; i < orders_cols_.size(); i++) {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch orders_batch,
+ ExecBatch::Make(std::move(orders_result)));
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ orders_output_queue_.emplace(std::move(orders_batch));
+ }
+ }
+ std::vector<ExecBatch> lineitem_results;
+ for (size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++) {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for (size_t icol = 0; icol < lineitem_cols_.size(); icol++) {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ lineitem_batches_generated_.fetch_add(static_cast<int64_t>(lineitem_results.size()));
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ for (size_t i = 1; i < lineitem_results.size(); i++)
+ lineitem_output_queue_.emplace(std::move(lineitem_results[i]));
+ }
+ return std::move(lineitem_results[0]);
+ }
+
+ private:
+#define FOR_EACH_ORDERS_COLUMN(F) \
+ F(O_ORDERKEY) \
+ F(O_CUSTKEY) \
+ F(O_ORDERSTATUS) \
+ F(O_TOTALPRICE) \
+ F(O_ORDERDATE) \
+ F(O_ORDERPRIORITY) \
+ F(O_CLERK) \
+ F(O_SHIPPRIORITY) \
+ F(O_COMMENT)
+
+#define FOR_EACH_LINEITEM_COLUMN(F) \
+ F(L_ORDERKEY) \
+ F(L_PARTKEY) \
+ F(L_SUPPKEY) \
+ F(L_LINENUMBER) \
+ F(L_QUANTITY) \
+ F(L_EXTENDEDPRICE) \
+ F(L_DISCOUNT) \
+ F(L_TAX) \
+ F(L_RETURNFLAG) \
+ F(L_LINESTATUS) \
+ F(L_SHIPDATE) \
+ F(L_COMMITDATE) \
+ F(L_RECEIPTDATE) \
+ F(L_SHIPINSTRUCT) \
+ F(L_SHIPMODE) \
+ F(L_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct ORDERS {
+ enum {
+ FOR_EACH_ORDERS_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct LINEITEM {
+ enum {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, ORDERS::col},
+ const std::unordered_map<std::string, int> kOrdersNameMap = {
+ FOR_EACH_ORDERS_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, LINEITEM::col},
+ const std::unordered_map<std::string, int> kLineitemNameMap = {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ const std::vector<GenerateColumnFn> kOrdersGenerators = {
+ FOR_EACH_ORDERS_COLUMN(MAKE_FN_ARRAY)};
+ const std::vector<GenerateColumnFn> kLineitemGenerators = {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kOrdersTypes = {int32(),
+ int32(),
+ fixed_size_binary(1),
+ decimal(12, 2),
+ date32(),
+ fixed_size_binary(15),
+ fixed_size_binary(15),
+ int32(),
+ utf8()};
+
+ const std::vector<std::shared_ptr<DataType>> kLineitemTypes = {
+ int32(),
+ int32(),
+ int32(),
+ int32(),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ fixed_size_binary(1),
+ fixed_size_binary(1),
+ date32(),
+ date32(),
+ date32(),
+ fixed_size_binary(25),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status AllocateOrdersBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.orders[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.orders_to_generate * byte_width));
+ ArrayData ad(kOrdersTypes[column], tld.orders_to_generate,
+ {nullptr, std::move(buff)});
+ tld.orders[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status O_ORDERKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERKEY));
+ int32_t* o_orderkey = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ int32_t orderkey_index = static_cast<int32_t>(tld.orderkey_start + i);
+ int32_t index_of_run = orderkey_index / 8;
+ int32_t index_in_run = orderkey_index % 8;
+ o_orderkey[i] = (index_of_run * 32 + index_in_run + 1);
+ ARROW_DCHECK(1 <= o_orderkey[i] && o_orderkey[i] <= 4 * orders_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CUSTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_CUSTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CUSTKEY));
+
+ // Spec says it must be a random number between 1 and SF*150000 that is not
+ // divisible by 3. Rather than repeatedly generating numbers until we get to
+ // a non-divisible-by-3 number, we just generate a number between
+ // 0 and SF * 50000 - 1, multiply by 3, and then add either 1 or 2.
+ int32_t sf_50k = static_cast<int32_t>(scale_factor_ * 50000);
+ std::uniform_int_distribution<int32_t> base_dist(0, sf_50k - 1);
+ std::uniform_int_distribution<int32_t> offset_dist(1, 2);
+ int32_t* o_custkey = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_CUSTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++)
+ o_custkey[i] = 3 * base_dist(tld.rng) + offset_dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERSTATUS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERSTATUS].kind() == Datum::NONE) {
+ RETURN_NOT_OK(L_LINESTATUS(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERSTATUS));
+
+ char* o_orderstatus = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_ORDERSTATUS].array()->buffers[1]->mutable_data());
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ bool all_f = true;
+ bool all_o = true;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ const char* l_linestatus = reinterpret_cast<const char*>(
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->buffers[1]->data());
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++, batch_offset++) {
+ all_f &= l_linestatus[batch_offset] == 'F';
+ all_o &= l_linestatus[batch_offset] == 'O';
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ ARROW_DCHECK(!(all_f && all_o));
+ if (all_f)
+ o_orderstatus[iorder] = 'F';
+ else if (all_o)
+ o_orderstatus[iorder] = 'O';
+ else
+ o_orderstatus[iorder] = 'P';
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_TOTALPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_TOTALPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(L_EXTENDEDPRICE(thread_index));
+ RETURN_NOT_OK(L_TAX(thread_index));
+ RETURN_NOT_OK(L_DISCOUNT(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_TOTALPRICE));
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ int64_t sum = 0;
+ Decimal128* o_totalprice = reinterpret_cast<Decimal128*>(
+ tld.orders[ORDERS::O_TOTALPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ const Decimal128* l_extendedprice = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->buffers[1]->data());
+ const Decimal128* l_tax = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->data());
+ const Decimal128* l_discount = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->buffers[1]->data());
+
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++, batch_offset++) {
+ int64_t eprice = static_cast<int64_t>(l_extendedprice[batch_offset]);
+ int64_t tax = static_cast<int64_t>(l_tax[batch_offset]);
+ int64_t discount = static_cast<int64_t>(l_discount[batch_offset]);
+ sum += (eprice * (100 + tax) * (100 - discount));
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ sum /= 100 * 100;
+ o_totalprice[iorder] = {sum};
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERDATE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERDATE));
+
+ std::uniform_int_distribution<uint32_t> dist(kStartDate, kEndDate - 151);
+ uint32_t* o_orderdate = reinterpret_cast<uint32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) o_orderdate[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERPRIORITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERPRIORITY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERPRIORITY));
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_ORDERPRIORITY]);
+ std::uniform_int_distribution<int32_t> dist(0, kNumPriorities - 1);
+ char* o_orderpriority = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_ORDERPRIORITY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ const char* str = Priorities[dist(tld.rng)];
+ std::strncpy(o_orderpriority + i * byte_width, str, byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CLERK(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_CLERK].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CLERK));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_CLERK]);
+ int64_t max_clerk_id = static_cast<int64_t>(scale_factor_ * 1000);
+ std::uniform_int_distribution<int64_t> dist(1, max_clerk_id);
+ char* o_clerk = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_CLERK].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ const char* clerk = "Clerk#";
+ const size_t clerk_length = std::strlen(clerk);
+ int64_t clerk_number = dist(tld.rng);
+ char* output = o_clerk + i * byte_width;
+ std::strncpy(output, clerk, byte_width);
+ AppendNumberPaddedToNineDigits(output + clerk_length, clerk_number);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_SHIPPRIORITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_SHIPPRIORITY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_SHIPPRIORITY));
+ int32_t* o_shippriority = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_SHIPPRIORITY].array()->buffers[1]->mutable_data());
+ std::memset(o_shippriority, 0, tld.orders_to_generate * sizeof(int32_t));
+ }
+ return Status::OK();
+ }
+
+ Status O_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.orders[ORDERS::O_COMMENT],
+ g_text.GenerateComments(tld.orders_to_generate, 19, 78, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ Status GenerateRowCounts(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ std::uniform_int_distribution<int> length_dist(1, 7);
+ tld.lineitem_to_generate = 0;
+ tld.items_per_order.clear();
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ int length = length_dist(tld.rng);
+ tld.items_per_order.push_back(length);
+ tld.lineitem_to_generate += length;
+ }
+ int64_t num_batches =
+ (tld.first_batch_offset + tld.lineitem_to_generate + batch_size_ - 1) /
+ batch_size_;
+ tld.lineitem.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.lineitem) {
+ batch.resize(LINEITEM::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocateLineItemBufferIfNeeded(size_t thread_index, size_t ibatch, int column,
+ size_t& out_batch_offset) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.lineitem[ibatch][column].kind() == Datum::NONE) {
+ int32_t byte_width = arrow::internal::GetByteWidth(*kLineitemTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kLineitemTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.lineitem[ibatch][column] = std::move(ad);
+ out_batch_offset = 0;
+ }
+ if (ibatch == 0) out_batch_offset = tld.first_batch_offset;
+
+ return Status::OK();
+ }
+
+ Status L_ORDERKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_ORDERKEY]) {
+ tld.generated_lineitem[LINEITEM::L_ORDERKEY] = true;
+ RETURN_NOT_OK(O_ORDERKEY(thread_index));
+ const int32_t* o_orderkey = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_ORDERKEY, batch_offset));
+ int32_t* l_linenumber =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_ORDERKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_linenumber[batch_offset++] = o_orderkey[iorder];
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_ORDERKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_PARTKEY]) {
+ tld.generated_lineitem[LINEITEM::L_PARTKEY] = true;
+
+ size_t ibatch = 0;
+ int32_t max_partkey = static_cast<int32_t>(scale_factor_ * 200000);
+ std::uniform_int_distribution<int32_t> dist(1, max_partkey);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_PARTKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_partkey =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_partkey[batch_offset] = dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SUPPKEY]) {
+ tld.generated_lineitem[LINEITEM::L_SUPPKEY] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 3);
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset = 0;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SUPPKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_suppkey =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const int32_t* l_partkey = reinterpret_cast<const int32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ for (int64_t i = 0; i < next_run; i++) {
+ int32_t supplier = dist(tld.rng);
+ int32_t partkey = l_partkey[batch_offset];
+ // Fun fact: the parentheses for this expression are unbalanced in the TPC-H
+ // spec.
+ l_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SUPPKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINENUMBER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_LINENUMBER]) {
+ tld.generated_lineitem[LINEITEM::L_LINENUMBER] = true;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_LINENUMBER, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_linenumber =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_LINENUMBER]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++) {
+ l_linenumber[batch_offset++] = (iline + 1);
+ ARROW_DCHECK(1 <= (iline + 1) && (iline + 1) <= 7);
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINENUMBER].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_QUANTITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_QUANTITY]) {
+ tld.generated_lineitem[LINEITEM::L_QUANTITY] = true;
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int64_t> dist(1, 50);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_QUANTITY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_quantity =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_QUANTITY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t i = 0; i < next_run; i++) {
+ // Multiply by 100 because the type is decimal(12, 2), so the decimal goes after
+ // two digits
+ int64_t quantity = dist(tld.rng) * 100;
+ l_quantity[batch_offset++] = {quantity};
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_EXTENDEDPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE]) {
+ tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+ RETURN_NOT_OK(L_QUANTITY(thread_index));
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_EXTENDEDPRICE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ const int32_t* l_partkey = reinterpret_cast<const int32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ const Decimal128* l_quantity = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->buffers[1]->data());
+ Decimal128* l_extendedprice =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ int64_t partkey = static_cast<int64_t>(l_partkey[batch_offset]);
+ // Divide by 100 to recover the integer representation (not Decimal).
+ int64_t quantity = static_cast<int64_t>(l_quantity[batch_offset]) / 100;
+
+ // Spec says to divide by 100, but that happens automatically due to this being
+ // stored to two decimal points.
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ int64_t extended_price = retail_price * quantity;
+ l_extendedprice[batch_offset] = {extended_price};
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_DISCOUNT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_DISCOUNT]) {
+ tld.generated_lineitem[LINEITEM::L_DISCOUNT] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 10);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_DISCOUNT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_discount =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_DISCOUNT]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_discount[batch_offset] = {dist(tld.rng)};
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_TAX(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_TAX]) {
+ tld.generated_lineitem[LINEITEM::L_TAX] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 8);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_TAX, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_tax = reinterpret_cast<Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_tax[batch_offset] = {dist(tld.rng)};
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RETURNFLAG(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_RETURNFLAG]) {
+ tld.generated_lineitem[LINEITEM::L_RETURNFLAG] = true;
+ RETURN_NOT_OK(L_RECEIPTDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<uint32_t> dist;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_RETURNFLAG, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_returnflag =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_receiptdate = reinterpret_cast<const uint32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ if (l_receiptdate[batch_offset] <= kCurrentDate) {
+ uint32_t r = dist(tld.rng);
+ l_returnflag[batch_offset] = (r % 2 == 1) ? 'R' : 'A';
+ } else {
+ l_returnflag[batch_offset] = 'N';
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINESTATUS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_LINESTATUS]) {
+ tld.generated_lineitem[LINEITEM::L_LINESTATUS] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_LINESTATUS, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_linestatus =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_LINESTATUS]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_shipdate =
+ reinterpret_cast<const uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ if (l_shipdate[batch_offset] > kCurrentDate)
+ l_linestatus[batch_offset] = 'O';
+ else
+ l_linestatus[batch_offset] = 'F';
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPDATE]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPDATE] = true;
+ RETURN_NOT_OK(O_ORDERDATE(thread_index));
+ const int32_t* o_orderdate = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(1, 121);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SHIPDATE, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_shipdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_shipdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMITDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_COMMITDATE]) {
+ tld.generated_lineitem[LINEITEM::L_COMMITDATE] = true;
+ const int32_t* o_orderdate = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(30, 90);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_COMMITDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_commitdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_COMMITDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_commitdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMITDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RECEIPTDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_RECEIPTDATE]) {
+ tld.generated_lineitem[LINEITEM::L_RECEIPTDATE] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(1, 30);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_RECEIPTDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_receiptdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_shipdate =
+ reinterpret_cast<const uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_receiptdate[batch_offset] = l_shipdate[batch_offset] + dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPINSTRUCT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT] = true;
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPINSTRUCT]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumInstructions - 1);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_SHIPINSTRUCT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_shipinstruct =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ const char* str = Instructions[dist(tld.rng)];
+ // Note that we don't have to memset the buffer to 0 because strncpy pads each
+ // string with 0's anyway
+ std::strncpy(l_shipinstruct + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPMODE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPMODE]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPMODE] = true;
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPMODE]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumModes - 1);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SHIPMODE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_shipmode =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_SHIPMODE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ const char* str = Modes[dist(tld.rng)];
+ std::strncpy(l_shipmode + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPMODE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_COMMENT]) {
+ tld.generated_lineitem[LINEITEM::L_COMMENT] = true;
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ // Comments are kind of sneaky: we always generate the full batch and then just
+ // bump the length
+ if (tld.lineitem[ibatch][LINEITEM::L_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.lineitem[ibatch][LINEITEM::L_COMMENT],
+ g_text.GenerateComments(batch_size_, 10, 43, tld.rng));
+ batch_offset = 0;
+ }
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ batch_offset += next_run;
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMENT].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ std::vector<Datum> orders;
+ int64_t orders_to_generate;
+ int64_t orderkey_start;
+
+ std::vector<std::vector<Datum>> lineitem;
+ std::vector<int> items_per_order;
+ int64_t lineitem_to_generate;
+ int64_t first_batch_offset;
+ std::bitset<LINEITEM::kNumCols> generated_lineitem;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex orders_output_queue_mutex_;
+ std::mutex lineitem_output_queue_mutex_;
+ std::queue<ExecBatch> orders_output_queue_;
+ std::queue<ExecBatch> lineitem_output_queue_;
+ int64_t batch_size_;
+ double scale_factor_;
+ int64_t orders_rows_to_generate_;
+ int64_t orders_rows_generated_;
+ std::vector<int> orders_cols_;
+ std::vector<int> lineitem_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<size_t> orders_batches_generated_ = {0};
+ std::atomic<size_t> lineitem_batches_generated_ = {0};
+};
+
+class SupplierGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, double scale_factor, int64_t batch_size,
+ int64_t seed) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 10000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, gen_list_));
+
+ random::pcg64_fast seed_rng(seed);
+ seed_ = kSeedDist(seed_rng);
+
+ random::pcg32_fast rng(kSeedDist(seed_rng));
+ std::uniform_int_distribution<int64_t> dist(0, rows_to_generate_ - 1);
+ size_t num_special_rows = static_cast<size_t>(5 * scale_factor_);
+ std::unordered_set<int64_t> good_rows_set;
+ while (good_rows_set.size() < num_special_rows) {
+ int64_t row = dist(rng);
+ good_rows_set.insert(row);
+ }
+ std::unordered_set<int64_t> bad_rows_set;
+ while (bad_rows_set.size() < num_special_rows) {
+ int64_t bad_row;
+ do {
+ bad_row = dist(rng);
+ } while (good_rows_set.find(bad_row) != good_rows_set.end());
+ bad_rows_set.insert(bad_row);
+ }
+ good_rows_.clear();
+ bad_rows_.clear();
+ good_rows_.insert(good_rows_.end(), good_rows_set.begin(), good_rows_set.end());
+ bad_rows_.insert(bad_rows_.end(), bad_rows_set.begin(), bad_rows_set.end());
+ std::sort(good_rows_.begin(), good_rows_.end());
+ std::sort(bad_rows_.begin(), bad_rows_.end());
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) tld.rng.seed(GetRandomSeed());
+
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(S_SUPPKEY) \
+ F(S_NAME) \
+ F(S_ADDRESS) \
+ F(S_NATIONKEY) \
+ F(S_PHONE) \
+ F(S_ACCTBAL) \
+ F(S_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct SUPPLIER {
+ enum {
+ FOR_EACH_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) {#col, SUPPLIER::col},
+ const std::unordered_map<std::string, int> kNameMap = {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kGenerators = {FOR_EACH_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(), fixed_size_binary(25), utf8(),
+ int32(), fixed_size_binary(15), decimal(12, 2),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index) {
+ if (done_.load()) return Status::OK();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.suppkey_start = rows_generated_.fetch_add(batch_size_);
+ if (tld.suppkey_start >= rows_to_generate_) return Status::OK();
+
+ tld.to_generate = std::min(batch_size_, rows_to_generate_ - tld.suppkey_start);
+
+ tld.batch.resize(SUPPLIER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for (int col : gen_list_) RETURN_NOT_OK(kGenerators[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for (size_t i = 0; i < gen_list_.size(); i++) {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_outputted_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_outputted_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if (is_last_batch) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ return Status::OK();
+ }
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(kTypes[column], tld.to_generate, {nullptr, std::move(buff)});
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status S_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_SUPPKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_SUPPKEY));
+ int32_t* s_suppkey = reinterpret_cast<int32_t*>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ s_suppkey[irow] = static_cast<int32_t>(tld.suppkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_NAME].kind() == Datum::NONE) {
+ RETURN_NOT_OK(S_SUPPKEY(thread_index));
+ const int32_t* s_suppkey = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->data());
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NAME));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[SUPPLIER::S_NAME]);
+ char* s_name = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_NAME].array()->buffers[1]->mutable_data());
+ // Look man, I'm just following the spec ok? Section 4.2.3 as of March 1 2022
+ const char* supplier = "Supplie#r";
+ const size_t supplier_length = std::strlen(supplier);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = s_name + byte_width * irow;
+ std::strncpy(out, supplier, byte_width);
+ AppendNumberPaddedToNineDigits(out + supplier_length, s_suppkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ADDRESS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_ADDRESS].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[SUPPLIER::S_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status S_NATIONKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_NATIONKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t* s_nationkey = reinterpret_cast<int32_t*>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status S_PHONE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_PHONE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(S_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[SUPPLIER::S_PHONE]);
+ const int32_t* s_nationkey = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->data());
+ char* s_phone = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_PHONE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ GeneratePhoneNumber(s_phone + irow * byte_width, tld.rng, s_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ACCTBAL(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_ACCTBAL].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_ACCTBAL));
+ Decimal128* s_acctbal = reinterpret_cast<Decimal128*>(
+ tld.batch[SUPPLIER::S_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_acctbal[irow] = {dist(tld.rng)};
+ }
+ return Status::OK();
+ }
+
+ Status S_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[SUPPLIER::S_COMMENT],
+ g_text.GenerateComments(tld.to_generate, 25, 100, tld.rng));
+ ModifyComments(thread_index, "Recommends", good_rows_);
+ ModifyComments(thread_index, "Complaints", bad_rows_);
+ }
+ return Status::OK();
+ }
+
+ void ModifyComments(size_t thread_index, const char* review,
+ const std::vector<int64_t>& indices) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ const int32_t* offsets = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[1]->data());
+ char* str = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[2]->mutable_data());
+ const char* customer = "Customer";
+ const int32_t customer_length = static_cast<int32_t>(std::strlen(customer));
+ const int32_t review_length = static_cast<int32_t>(std::strlen(review));
+
+ auto it = std::lower_bound(indices.begin(), indices.end(), tld.suppkey_start);
+ for (; it != indices.end() && *it < tld.suppkey_start + tld.to_generate; it++) {
+ int64_t idx_in_batch = *it - tld.suppkey_start;
+ char* out = str + offsets[idx_in_batch];
+ int32_t str_length = offsets[idx_in_batch + 1] - offsets[idx_in_batch];
+ std::uniform_int_distribution<int32_t> gap_dist(
+ 0, str_length - customer_length - review_length);
+ int32_t gap = gap_dist(tld.rng);
+ int32_t total_length = customer_length + gap + review_length;
+ std::uniform_int_distribution<int32_t> start_dist(0, str_length - total_length);
+ int32_t start = start_dist(tld.rng);
+ std::memcpy(out + start, customer, customer_length);
+ std::memcpy(out + start + customer_length + gap, review, review_length);
+ }
+ }
+
+ struct ThreadLocalData {
+ random::pcg32_fast rng;
+ int64_t suppkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+ std::vector<int64_t> good_rows_;
+ std::vector<int64_t> bad_rows_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_;
+ std::atomic<int64_t> rows_generated_;
+ double scale_factor_;
+ int64_t batch_size_;
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class PartGenerator : public TpchTableGenerator {
+ public:
+ explicit PartGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, double scale_factor, int64_t batch_size,
+ int64_t seed) override {
+ seed_ = seed;
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetPartOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_, seed_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch, gen_->NextPartBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->part_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ double scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class PartSuppGenerator : public TpchTableGenerator {
+ public:
+ explicit PartSuppGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, double scale_factor, int64_t batch_size,
+ int64_t seed) override {
+ seed_ = seed;
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetPartSuppOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_, seed_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextPartSuppBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->partsupp_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ double scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class CustomerGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, double scale_factor, int64_t batch_size,
+ int64_t seed) override {
+ seed_ = seed;
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, gen_list_));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ thread_local_data_.resize(num_threads);
+ random::pcg64_fast seed_rng(seed_);
+ for (ThreadLocalData& tld : thread_local_data_) tld.rng.seed(kSeedDist(seed_rng));
+
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(C_CUSTKEY) \
+ F(C_NAME) \
+ F(C_ADDRESS) \
+ F(C_NATIONKEY) \
+ F(C_PHONE) \
+ F(C_ACCTBAL) \
+ F(C_MKTSEGMENT) \
+ F(C_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct CUSTOMER {
+ enum {
+ FOR_EACH_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) {#col, CUSTOMER::col},
+ const std::unordered_map<std::string, int> kNameMap = {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kGenerators = {FOR_EACH_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(),
+ utf8(),
+ utf8(),
+ int32(),
+ fixed_size_binary(15),
+ decimal(12, 2),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index) {
+ if (done_.load()) return Status::OK();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.custkey_start = rows_generated_.fetch_add(batch_size_);
+ if (tld.custkey_start >= rows_to_generate_) return Status::OK();
+
+ tld.to_generate = std::min(batch_size_, rows_to_generate_ - tld.custkey_start);
+
+ tld.batch.resize(CUSTOMER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for (int col : gen_list_) RETURN_NOT_OK(kGenerators[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for (size_t i = 0; i < gen_list_.size(); i++) {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_generated_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_generated_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if (is_last_batch) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ return Status::OK();
+ }
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(kTypes[column], tld.to_generate, {nullptr, std::move(buff)});
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status C_CUSTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_CUSTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_CUSTKEY));
+ int32_t* c_custkey = reinterpret_cast<int32_t*>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ c_custkey[irow] = static_cast<int32_t>(tld.custkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_NAME].kind() == Datum::NONE) {
+ RETURN_NOT_OK(C_CUSTKEY(thread_index));
+ const int32_t* c_custkey = reinterpret_cast<const int32_t*>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->data());
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ const char* customer = "Customer#";
+ const size_t customer_length = std::strlen(customer);
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ int num_digits = GetNumDigits(c_custkey[irow]);
+ int num_chars = std::max(num_digits, 9);
+ offsets[irow + 1] =
+ static_cast<int32_t>(offsets[irow] + num_chars + customer_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[tld.to_generate]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = str + offsets[irow];
+ std::memcpy(out, customer, customer_length);
+ AppendNumberPaddedToNineDigits(out + customer_length, c_custkey[irow]);
+ }
+ ArrayData ad(utf8(), tld.to_generate,
+ {nullptr, std::move(offset_buff), std::move(str_buff)});
+ tld.batch[CUSTOMER::C_NAME] = std::move(ad);
+ }
+ return Status::OK();
+ }
+
+ Status C_ADDRESS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_ADDRESS].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[CUSTOMER::C_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status C_NATIONKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_NATIONKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t* c_nationkey = reinterpret_cast<int32_t*>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status C_PHONE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_PHONE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(C_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[CUSTOMER::C_PHONE]);
+ const int32_t* c_nationkey = reinterpret_cast<const int32_t*>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->data());
+ char* c_phone = reinterpret_cast<char*>(
+ tld.batch[CUSTOMER::C_PHONE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ GeneratePhoneNumber(c_phone + irow * byte_width, tld.rng, c_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_ACCTBAL(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_ACCTBAL].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_ACCTBAL));
+ Decimal128* c_acctbal = reinterpret_cast<Decimal128*>(
+ tld.batch[CUSTOMER::C_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_acctbal[irow] = {dist(tld.rng)};
+ }
+ return Status::OK();
+ }
+
+ Status C_MKTSEGMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_MKTSEGMENT].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_MKTSEGMENT));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[CUSTOMER::C_MKTSEGMENT]);
+ char* c_mktsegment = reinterpret_cast<char*>(
+ tld.batch[CUSTOMER::C_MKTSEGMENT].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int32_t> dist(0, kNumSegments - 1);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = c_mktsegment + irow * byte_width;
+ int str_idx = dist(tld.rng);
+ std::strncpy(out, Segments[str_idx], byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[CUSTOMER::C_COMMENT],
+ g_text.GenerateComments(tld.to_generate, 29, 116, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ random::pcg32_fast rng;
+ int64_t custkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_{0};
+ std::atomic<int64_t> rows_generated_ = {0};
+ double scale_factor_{0};
+ int64_t batch_size_{0};
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class OrdersGenerator : public TpchTableGenerator {
+ public:
+ explicit OrdersGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, double scale_factor, int64_t batch_size,
+ int64_t seed) override {
+ seed_ = seed;
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetOrdersOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_, seed_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch, gen_->NextOrdersBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->orders_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ double scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class LineitemGenerator : public TpchTableGenerator {
+ public:
+ explicit LineitemGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, double scale_factor, int64_t batch_size,
+ int64_t seed) override {
+ seed_ = seed;
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetLineItemOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_, seed_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextLineItemBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->lineitem_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ double scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class NationGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, double /*scale_factor*/,
+ int64_t /*batch_size*/, int64_t seed) override {
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, column_indices_));
+ seed_ = seed;
+ rng_.seed(seed_);
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t /*num_threads*/, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback /*schedule_task_callback*/) override {
+ std::shared_ptr<Buffer> N_NATIONKEY_buffer =
+ Buffer::Wrap(kNationKey, sizeof(kNationKey));
+ ArrayData N_NATIONKEY_arraydata(int32(), kRowCount,
+ {nullptr, std::move(N_NATIONKEY_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> N_NAME_buffer,
+ AllocateBuffer(kRowCount * kNameByteWidth));
+ char* N_NAME = reinterpret_cast<char*>(N_NAME_buffer->mutable_data());
+ for (size_t i = 0; i < kRowCount; i++)
+ std::strncpy(N_NAME + kNameByteWidth * i, kCountryNames[i], kNameByteWidth);
+ ArrayData N_NAME_arraydata(fixed_size_binary(kNameByteWidth), kRowCount,
+ {nullptr, std::move(N_NAME_buffer)});
+
+ std::shared_ptr<Buffer> N_REGIONKEY_buffer =
+ Buffer::Wrap(kRegionKey, sizeof(kRegionKey));
+ ArrayData N_REGIONKEY_arraydata(int32(), kRowCount,
+ {nullptr, std::move(N_REGIONKEY_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(Datum N_COMMENT_datum,
+ g_text.GenerateComments(kRowCount, 31, 114, rng_));
+
+ std::vector<Datum> fields = {
+ std::move(N_NATIONKEY_arraydata), std::move(N_NAME_arraydata),
+ std::move(N_REGIONKEY_arraydata), std::move(N_COMMENT_datum)};
+
+ std::vector<Datum> result;
+ for (const int& col : column_indices_) result.push_back(fields[col]);
+ ARROW_ASSIGN_OR_RAISE(ExecBatch batch, ExecBatch::Make(std::move(result)));
+ output_callback(std::move(batch));
+ finished_callback(static_cast<int64_t>(1));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ random::pcg32_fast rng_;
+
+ static constexpr size_t kRowCount = 25;
+ static constexpr int32_t kNameByteWidth = 25;
+ static constexpr int32_t kNationKey[kRowCount] = {0, 1, 2, 3, 4, 5, 6, 7, 8,
+ 9, 10, 11, 12, 13, 14, 15, 16, 17,
+ 18, 19, 20, 21, 22, 23, 24};
+ static constexpr const char* kCountryNames[kRowCount] = {
+ "ALGERIA", "ARGENTINA", "BRAZIL", "CANADA", "EGYPT",
+ "ETHIOPIA", "FRANCE", "GERMANY", "INDIA", "INDONESIA",
+ "IRAN", "IRAQ", "JAPAN", "JORDAN", "KENYA",
+ "MOROCCO", "MOZAMBIQUE", "PERU", "CHINA", "ROMANIA",
+ "SAUDI ARABIA", "VIETNAM", "RUSSIA", "UNITED KINGDOM", "UNITED STATES"};
+ static constexpr int32_t kRegionKey[kRowCount] = {0, 1, 1, 1, 4, 0, 3, 3, 2, 2, 4, 4, 2,
+ 4, 0, 0, 0, 1, 2, 3, 4, 2, 3, 3, 1};
+
+ struct NATION {
+ enum {
+ N_NATIONKEY,
+ N_NAME,
+ N_REGIONKEY,
+ N_COMMENT,
+ };
+ };
+
+ const std::unordered_map<std::string, int> kNameMap = {
+ {"N_NATIONKEY", NATION::N_NATIONKEY},
+ {"N_NAME", NATION::N_NAME},
+ {"N_REGIONKEY", NATION::N_REGIONKEY},
+ {"N_COMMENT", NATION::N_COMMENT},
+ };
+
+ std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(),
+ fixed_size_binary(kNameByteWidth),
+ int32(),
+ utf8(),
+ };
+
+ std::shared_ptr<Schema> schema_;
+ std::vector<int> column_indices_;
+};
+
+constexpr int32_t NationGenerator::kNationKey[NationGenerator::kRowCount];
+constexpr const char* NationGenerator::kCountryNames[NationGenerator::kRowCount];
+constexpr int32_t NationGenerator::kRegionKey[NationGenerator::kRowCount];
+
+class RegionGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, double /*scale_factor*/,
+ int64_t /*batch_size*/, int64_t seed) override {
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, column_indices_));
+ seed_ = seed;
+ rng_.seed(seed_);
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback /*schedule_task_callback*/) override {
+ std::shared_ptr<Buffer> R_REGIONKEY_buffer =
+ Buffer::Wrap(kRegionKey, sizeof(kRegionKey));
+ ArrayData R_REGIONKEY_arraydata(int32(), kRowCount,
+ {nullptr, std::move(R_REGIONKEY_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> R_NAME_buffer,
+ AllocateBuffer(kRowCount * kNameByteWidth));
+ char* R_NAME_data = reinterpret_cast<char*>(R_NAME_buffer->mutable_data());
+ for (size_t i = 0; i < kRowCount; i++)
+ std::strncpy(R_NAME_data + kNameByteWidth * i, kRegionNames[i], kNameByteWidth);
+ ArrayData R_NAME_arraydata(kTypes[static_cast<int>(REGION::R_NAME)], kRowCount,
+ {nullptr, std::move(R_NAME_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(Datum R_COMMENT_datum,
+ g_text.GenerateComments(kRowCount, 31, 115, rng_));
+
+ std::vector<Datum> fields = {std::move(R_REGIONKEY_arraydata),
+ std::move(R_NAME_arraydata), std::move(R_COMMENT_datum)};
+ std::vector<Datum> result;
+ for (const int& col : column_indices_) result.push_back(fields[col]);
+ ARROW_ASSIGN_OR_RAISE(ExecBatch batch, ExecBatch::Make(std::move(result)));
+ output_callback(std::move(batch));
+ finished_callback(static_cast<int64_t>(1));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ random::pcg32_fast rng_;
+
+ static constexpr size_t kRowCount = 5;
+ static constexpr int32_t kNameByteWidth = 25;
+ static constexpr int32_t kRegionKey[kRowCount] = {0, 1, 2, 3, 4};
+ static constexpr const char* kRegionNames[kRowCount] = {"AFRICA", "AMERICA", "ASIA",
+ "EUROPE", "MIDDLE EAST"};
+
+ struct REGION {
+ enum {
+ R_REGIONKEY,
+ R_NAME,
+ R_COMMENT,
+ kNumColumns,
+ };
+ };
+
+ const std::unordered_map<std::string, int> kNameMap = {
+ {"R_REGIONKEY", REGION::R_REGIONKEY},
+ {"R_NAME", REGION::R_NAME},
+ {"R_COMMENT", REGION::R_COMMENT},
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(),
+ fixed_size_binary(kNameByteWidth),
+ utf8(),
+ };
+
+ std::shared_ptr<Schema> schema_;
+ std::vector<int> column_indices_;
+};
+
+constexpr int32_t RegionGenerator::kRegionKey[RegionGenerator::kRowCount];
+constexpr const char* RegionGenerator::kRegionNames[RegionGenerator::kRowCount];
+
+class TpchNode : public ExecNode {
+ public:
+ TpchNode(ExecPlan* plan, const char* name,
+ std::unique_ptr<TpchTableGenerator> generator)
+ : ExecNode(plan, {}, {}, generator->schema(), /*num_outputs=*/1),
+ name_(name),
+ generator_(std::move(generator)) {}
+
+ const char* kind_name() const override { return name_; }
+
+ [[noreturn]] static void NoInputs() {
+ Unreachable("TPC-H node should never have any inputs");
+ }
+
+ [[noreturn]] void InputReceived(ExecNode*, ExecBatch) override { NoInputs(); }
+
+ [[noreturn]] void ErrorReceived(ExecNode*, Status) override { NoInputs(); }
+
+ [[noreturn]] void InputFinished(ExecNode*, int) override { NoInputs(); }
+
+ Status StartProducing() override {
+ return generator_->StartProducing(
+ thread_indexer_.Capacity(),
+ [this](ExecBatch batch) { this->OutputBatchCallback(std::move(batch)); },
+ [this](int64_t num_batches) { this->FinishedCallback(num_batches); },
+ [this](std::function<Status(size_t)> func) -> Status {
+ return this->ScheduleTaskCallback(std::move(func));
+ });
+ }
+
+ void PauseProducing(ExecNode* output) override {}
+ void ResumeProducing(ExecNode* output) override {}
+
+ void StopProducing(ExecNode* output) override {
+ DCHECK_EQ(output, outputs_[0]);
+ StopProducing();
+ }
+
+ void StopProducing() override {
+ if (generator_->Abort()) std::ignore = task_group_.End();
+ }
+
+ Future<> finished() override { return task_group_.OnFinished(); }
+
+ private:
+ void OutputBatchCallback(ExecBatch batch) {
+ outputs_[0]->InputReceived(this, std::move(batch));
+ }
+
+ void FinishedCallback(int64_t total_num_batches) {
+ outputs_[0]->InputFinished(this, static_cast<int>(total_num_batches));
+ std::ignore = task_group_.End();
+ }
+
+ Status ScheduleTaskCallback(std::function<Status(size_t)> func) {
+ auto executor = plan_->exec_context()->executor();
+ if (executor) {
+ RETURN_NOT_OK(task_group_.AddTaskIfNotEnded([&] {
Review comment:
Can you add a comment here explaining why it's ok to ignore tasks. Maybe something like:
// The first part of a generation task is figuring out which rows to emit.
// When we reach the end of a table generator we may end up creating
// some extra final tasks. The first one of these will mark the task group
// ended and the remaining generation tasks can be safely ignored.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou edited a comment on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou edited a comment on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082144362
Hmm, there's an ASAN issue that needs to be looked at before merging this. I can do that tomorrow if you want.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r825009715
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string> &columns,
+ const std::vector<std::shared_ptr<DataType>> &types,
+ const std::unordered_map<std::string, int> &name_map,
+ std::vector<int> &gen_list)
+ {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if(columns.empty())
+ {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for(auto pair : name_map)
+ {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ }
+ else
+ {
+ for(const std::string &col : columns)
+ {
+ auto entry = name_map.find(col);
+ if(entry == name_map.end())
+ return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+ }
+
+ static TpchText g_text;
+
+ Status TpchText::InitIfNeeded(random::pcg32_fast &rng)
+ {
+ if(done_.load())
+ return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(!text_)
+ {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char *out = reinterpret_cast<char *>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+ while(done_.load() == false)
+ {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while(GenerateSentence(try_offset, rng, temp_buff))
+ known_valid_offset = try_offset;
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(done_.load())
+ return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ int64_t memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ std::memcpy(out + generated_offset_, temp_buff, memcpy_size);
+ generated_offset_ += memcpy_size;
+ if(generated_offset_ == kTextBytes)
+ done_.store(true);
+ }
+ }
+ return Status::OK();
+ }
+
+ Result<Datum> TpchText::GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng)
+ {
+ RETURN_NOT_OK(InitIfNeeded(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer, AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for(size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + length_dist(rng);
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer, AllocateBuffer(offsets[num_comments]));
+ char *comments = reinterpret_cast<char *>(comment_buffer->mutable_data());
+ for(size_t i = 0; i < num_comments; i++)
+ {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments, { nullptr, std::move(offset_buffer), std::move(comment_buffer) });
+ return std::move(ad);
+ }
+
+ Result<Datum> RandomVString(
+ random::pcg32_fast &rng,
+ int64_t num_rows,
+ int32_t min_length,
+ int32_t max_length)
+ {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff, AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for(int64_t i = 1; i <= num_rows; i++)
+ offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff, AllocateBuffer(offsets[num_rows]));
+ char *str = reinterpret_cast<char *>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for(int32_t i = 0; i < offsets[num_rows]; i++)
+ str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, { nullptr, std::move(offset_buff), std::move(str_buff) });
+ return std::move(ad);
+ }
+
+ void AppendNumber(char *&out, int num_digits, int32_t x)
+ {
+ out += (num_digits - 1);
+ while(x > 0)
+ {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+ }
+
+ void GeneratePhoneNumber(
+ char *out,
+ random::pcg32_fast &rng,
+ int32_t country)
+ {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+ }
+
+ static constexpr uint32_t STARTDATE = 8035; // January 1, 1992 is 8035 days after January 1, 1970
+ static constexpr uint32_t CURRENTDATE = 9298; // June 17, 1995 is 9298 days after January 1, 1970
+ static constexpr uint32_t ENDDATE = 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+ const char *NameParts[] =
+ {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque", "black", "blanched", "blue",
+ "blush", "brown", "burlywood", "burnished", "chartreuse", "chiffon", "chocolate", "coral",
+ "cornflower", "cornsilk", "cream", "cyan", "dark", "deep", "dim", "dodger", "drab", "firebrick",
+ "floral", "forest", "frosted", "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender", "lawn", "lemon", "light", "lime", "linen",
+ "magenta", "maroon", "medium", "metallic", "midnight", "mint", "misty", "moccasin", "navajo",
+ "navy", "olive", "orange", "orchid", "pale", "papaya", "peach", "peru", "pink", "plum", "powder",
+ "puff", "purple", "red", "rose", "rosy", "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel", "tan", "thistle", "tomato", "turquoise", "violet",
+ "wheat", "white", "yellow",
+ };
+ static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+ const char *Types_1[] =
+ {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+ };
+ static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+ const char *Types_2[] =
+ {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+ };
+ static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+ const char *Types_3[] =
+ {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+ };
+ static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+ const char *Containers_1[] =
+ {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+ };
+ static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+ const char *Containers_2[] =
+ {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+ };
+ static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+ const char *Segments[] =
+ {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+ };
+ static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+ const char *Priorities[] =
+ {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+ };
+ static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+ const char *Instructions[] =
+ {
+ "DELIVER IN PERSON", "COLLECT COD", "NONE", "TAKE BACK RETURN",
+ };
+ static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+ const char *Modes[] =
+ {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+ };
+ static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+ const char *Nouns[] =
+ {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ", "dependencies ", "excuses ",
+ "platelets ", "asymptotes ", "courts ", "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ", "forges ", "braids ", "hockey players ", "frays ",
+ "warhorses ", "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ", "waters ", "orbits ", "gifts ", "sheaves ",
+ "depths ", "sentiments ", "decoys ", "realms ", "pains ", "grouches ", "escapades ",
+ };
+ static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+ const char *Verbs[] =
+ {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ", "boost ", "affix ", "detect ", "integrate ",
+ "maintain ", "nod ", "was ", "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ", "poach ", "serve ", "run ", "dazzle ",
+ "snooze ", "doze ", "unwind ", "kindle ", "play ", "hang ", "believe ", "doubt ",
+ };
+ static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+ const char *Adjectives[] =
+ {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ", "quiet ", "ruthless ", "thin ",
+ "close ", "dogged ", "daring ", "brave ", "stealthy ", "permanent ", "enticing ", "idle ", "busy ",
+ "regular ", "final ", "ironic ", "even ", "bold ", "silent ",
+ };
+ static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+ const char *Adverbs[] =
+ {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ", "blithely ", "quickly ", "fluffily ",
+ "slowly ", "quietly ", "ruthlessly ", "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ", "ironically ", "evenly ", "boldly ",
+ "silently ",
+ };
+ static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+ const char *Prepositions[] =
+ {
+ "about ", "above ", "according to ", "across ", "after ", "against ", "along ", "alongside of ", "among ",
+ "around ", "at ", "atop ", "before ", "behind ", "beneath ", "beside ", "besides ", "between ", "beyond ",
+ "beyond ", "by ", "despite ", "during ", "except ", "for ", "from ", "in place of ", "inside ", "instead of ",
+ "into ", "near ", "of ", "on ", "outside ", "over ", "past ", "since ", "through ", "throughout ", "to ",
+ "toward ", "under ", "until ", "up ", "upon ", "without ", "with ", "within ",
+ };
+ static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+ const char *Auxiliaries[] =
+ {
+ "do ", "may ", "might ", "shall ", "will ", "would ", "can ", "could ", "should ", "ought to ", "must ",
+ "will have to ", "shall have to ", "could have to ", "should have to ", "must have to ", "need to ", "try to ",
+ };
+ static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+ const char *Terminators[] =
+ {
+ ".", ";", ":", "?", "!", "--",
+ };
+ static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+ bool TpchText::GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices)
+ {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char *word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if(offset + length > kChunkSize)
+ return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+ }
+
+ bool TpchText::GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+ }
+
+ bool TpchText::GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+ }
+
+ bool TpchText::GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+ }
+
+ bool TpchText::GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+ }
+
+ bool TpchText::GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+ }
+
+ bool TpchText::GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+ }
+
+ bool TpchText::GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if(result)
+ std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+ }
+
+ bool TpchText::GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char *comma_space = ", ";
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ GenerateAdverb(offset, rng, arr);
+ GenerateAdjective(offset, rng, arr);
+ GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ bool TpchText::GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ bool TpchText::GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ const char *the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+ }
+
+ bool TpchText::GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 5:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ using GenerateColumnFn = std::function<Status(size_t)>;
+ class PartAndPartSupplierGenerator
+ {
+ public:
+ Status Init(
+ size_t num_threads,
+ int64_t batch_size,
+ float scale_factor)
+ {
+ if(!inited_)
+ {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ thread_local_data_.resize(num_threads);
+ for(ThreadLocalData &tld : thread_local_data_)
+ {
+ // 5 is the maximum number of different strings we need to concatenate
+ tld.string_indices.resize(5 * batch_size_);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const
+ {
+ return part_batches_generated_.load();
+ }
+
+ int64_t partsupp_batches_generated() const
+ {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, part_types_, part_name_map_, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, partsupp_types_, partsupp_name_map_, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch()
+ {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ bool all_generated = part_rows_generated_ == part_rows_to_generate_;
Review comment:
You can't move that if above because we may have generated all rows, but the queue may not be empty, so we need to check the queue first, and only then return if we've generated everything. I did get rid of the variable though and just pulled it into the `else if`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r825006578
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
Review comment:
Yeah, `gen_list` is a list of indices of columns. So if we e.g. wanted to project the same column twice or change the order, we fill the gen list with the desired indices in the desired order
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r825000258
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
Review comment:
Oh gotcha - yep I like that
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r824949907
##########
File path: cpp/src/arrow/compute/kernels/vector_sort.cc
##########
@@ -893,7 +893,8 @@ class TableSorter {
TableSorter(ExecContext* ctx, uint64_t* indices_begin, uint64_t* indices_end,
const Table& table, const SortOptions& options)
- : ctx_(ctx),
Review comment:
Yes, it was segfaulting when it was resetting the error because it was trying to delete an uninitialized object (or something like that, I fixed it a while ago and don't remember the exact details).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r824989592
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
Review comment:
Yes, you're right. This was based on my thinking that only one thread did the actual initialization.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string> &columns,
+ const std::vector<std::shared_ptr<DataType>> &types,
+ const std::unordered_map<std::string, int> &name_map,
+ std::vector<int> &gen_list)
+ {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if(columns.empty())
+ {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for(auto pair : name_map)
+ {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ }
+ else
+ {
+ for(const std::string &col : columns)
+ {
+ auto entry = name_map.find(col);
+ if(entry == name_map.end())
+ return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+ }
+
+ static TpchText g_text;
+
+ Status TpchText::InitIfNeeded(random::pcg32_fast &rng)
+ {
+ if(done_.load())
+ return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(!text_)
Review comment:
Yes, you're right. This was based on my thinking that only one thread did the actual initialization.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r824986313
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string> &columns,
+ const std::vector<std::shared_ptr<DataType>> &types,
+ const std::unordered_map<std::string, int> &name_map,
+ std::vector<int> &gen_list)
+ {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if(columns.empty())
+ {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for(auto pair : name_map)
+ {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ }
+ else
+ {
+ for(const std::string &col : columns)
+ {
+ auto entry = name_map.find(col);
+ if(entry == name_map.end())
+ return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+ }
+
+ static TpchText g_text;
+
+ Status TpchText::InitIfNeeded(random::pcg32_fast &rng)
+ {
+ if(done_.load())
+ return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(!text_)
+ {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char *out = reinterpret_cast<char *>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+ while(done_.load() == false)
+ {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while(GenerateSentence(try_offset, rng, temp_buff))
+ known_valid_offset = try_offset;
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(done_.load())
+ return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ int64_t memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ std::memcpy(out + generated_offset_, temp_buff, memcpy_size);
+ generated_offset_ += memcpy_size;
+ if(generated_offset_ == kTextBytes)
+ done_.store(true);
+ }
Review comment:
```suggestion
{
std::lock_guard<std::mutex> lock(text_guard_);
if(done_.load())
return Status::OK();
int64_t bytes_remaining = kTextBytes - generated_offset_;
int64_t memcpy_size = std::min(known_valid_offset, bytes_remaining);
int64_t start = out + generated_offset_;
generated_offset_ += memcpy_size;
if(generated_offset_ == kTextBytes)
done_.store(true);
}
std::memcpy(start, temp_buff, memcpy_size);
```
I doubt this would have a significant improvement to performance but it might be slightly faster.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r826493331
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3850 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ const char *NameParts[] =
+ {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque", "black", "blanched", "blue",
+ "blush", "brown", "burlywood", "burnished", "chartreuse", "chiffon", "chocolate", "coral",
+ "cornflower", "cornsilk", "cream", "cyan", "dark", "deep", "dim", "dodger", "drab", "firebrick",
+ "floral", "forest", "frosted", "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender", "lawn", "lemon", "light", "lime", "linen",
+ "magenta", "maroon", "medium", "metallic", "midnight", "mint", "misty", "moccasin", "navajo",
+ "navy", "olive", "orange", "orchid", "pale", "papaya", "peach", "peru", "pink", "plum", "powder",
+ "puff", "purple", "red", "rose", "rosy", "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel", "tan", "thistle", "tomato", "turquoise", "violet",
+ "wheat", "white", "yellow",
+ };
+ static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+ const char *Types_1[] =
+ {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+ };
+ static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+ const char *Types_2[] =
+ {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+ };
+ static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+ const char *Types_3[] =
+ {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+ };
+ static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+ const char *Containers_1[] =
+ {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+ };
+ static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+ const char *Containers_2[] =
+ {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+ };
+ static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+ const char *Segments[] =
+ {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+ };
+ static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+ const char *Priorities[] =
+ {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+ };
+ static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+ const char *Instructions[] =
+ {
+ "DELIVER IN PERSON", "COLLECT COD", "NONE", "TAKE BACK RETURN",
+ };
+ static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+ const char *Modes[] =
+ {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+ };
+ static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+ const char *Nouns[] =
+ {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ", "dependencies ", "excuses ",
+ "platelets ", "asymptotes ", "courts ", "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ", "forges ", "braids ", "hockey players ", "frays ",
+ "warhorses ", "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ", "waters ", "orbits ", "gifts ", "sheaves ",
+ "depths ", "sentiments ", "decoys ", "realms ", "pains ", "grouches ", "escapades ",
+ };
+ static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+ const char *Verbs[] =
+ {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ", "boost ", "affix ", "detect ", "integrate ",
+ "maintain ", "nod ", "was ", "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ", "poach ", "serve ", "run ", "dazzle ",
+ "snooze ", "doze ", "unwind ", "kindle ", "play ", "hang ", "believe ", "doubt ",
+ };
+ static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+ const char *Adjectives[] =
+ {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ", "quiet ", "ruthless ", "thin ",
+ "close ", "dogged ", "daring ", "brave ", "stealthy ", "permanent ", "enticing ", "idle ", "busy ",
+ "regular ", "final ", "ironic ", "even ", "bold ", "silent ",
+ };
+ static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+ const char *Adverbs[] =
+ {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ", "blithely ", "quickly ", "fluffily ",
+ "slowly ", "quietly ", "ruthlessly ", "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ", "ironically ", "evenly ", "boldly ",
+ "silently ",
+ };
+ static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+ const char *Prepositions[] =
+ {
+ "about ", "above ", "according to ", "across ", "after ", "against ", "along ", "alongside of ", "among ",
+ "around ", "at ", "atop ", "before ", "behind ", "beneath ", "beside ", "besides ", "between ", "beyond ",
+ "beyond ", "by ", "despite ", "during ", "except ", "for ", "from ", "in place of ", "inside ", "instead of ",
+ "into ", "near ", "of ", "on ", "outside ", "over ", "past ", "since ", "through ", "throughout ", "to ",
+ "toward ", "under ", "until ", "up ", "upon ", "without ", "with ", "within ",
+ };
+ static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+ const char *Auxiliaries[] =
+ {
+ "do ", "may ", "might ", "shall ", "will ", "would ", "can ", "could ", "should ", "ought to ", "must ",
+ "will have to ", "shall have to ", "could have to ", "should have to ", "must have to ", "need to ", "try to ",
+ };
+ static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+ const char *Terminators[] =
+ {
+ ".", ";", ":", "?", "!", "--",
+ };
+ static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+
+ // The spec says to generate a 300 MB string according to a grammar. This is a
+ // concurrent implementation of the generator. Each thread generates the text in
+ // (up to) 8KB chunks of text. The generator maintains a cursor into the
+ // 300 MB buffer. After generating the chunk, the cursor is incremented
+ // to reserve space, and the chunk is memcpy-d in.
+ // This text is used to generate the COMMENT columns. To generate a comment, the spec
+ // says to pick a random length and a random offset into the 300 MB buffer (it does
+ // not specify it should be word/sentence aligned), and that slice of text becomes
+ // the comment.
+ class TpchPseudotext
+ {
+ public:
+ Status EnsureInitialized(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ static TpchPseudotext g_text;
+
+ Status TpchPseudotext::EnsureInitialized(random::pcg32_fast &rng)
+ {
+ if(done_.load())
+ return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(!text_)
+ {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char *out = reinterpret_cast<char *>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while(!done_.load())
+ {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while(GenerateSentence(try_offset, rng, temp_buff))
+ known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(done_.load())
+ return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if(last_one)
+ done_.store(true);
+ }
+ return Status::OK();
+ }
+
+ Result<Datum> TpchPseudotext::GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng)
+ {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer, AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for(size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer, AllocateBuffer(offsets[num_comments]));
+ char *comments = reinterpret_cast<char *>(comment_buffer->mutable_data());
+ for(size_t i = 0; i < num_comments; i++)
+ {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments, { nullptr, std::move(offset_buffer), std::move(comment_buffer) });
+ return std::move(ad);
+ }
+
+ bool TpchPseudotext::GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices)
+ {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char *word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if(offset + length > kChunkSize)
+ return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+ }
+
+ bool TpchPseudotext::GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+ }
+
+ bool TpchPseudotext::GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+ }
+
+ bool TpchPseudotext::GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+ }
+
+ bool TpchPseudotext::GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+ }
+
+ bool TpchPseudotext::GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+ }
+
+ bool TpchPseudotext::GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+ }
+
+ bool TpchPseudotext::GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if(result)
+ std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+ }
+
+ bool TpchPseudotext::GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char *comma_space = ", ";
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ bool TpchPseudotext::GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ const char *the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+ }
+
+ bool TpchPseudotext::GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort()
+ {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string> &columns,
+ const std::vector<std::shared_ptr<DataType>> &types,
+ const std::unordered_map<std::string, int> &name_map,
+ std::vector<int> &gen_list)
+ {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if(columns.empty())
+ {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for(auto pair : name_map)
+ {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ }
+ else
+ {
+ for(const std::string &col : columns)
+ {
+ auto entry = name_map.find(col);
+ if(entry == name_map.end())
+ return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+ }
+
+ Result<Datum> RandomVString(
+ random::pcg32_fast &rng,
+ int64_t num_rows,
+ int32_t min_length,
+ int32_t max_length)
+ {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff, AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for(int64_t i = 1; i <= num_rows; i++)
+ offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff, AllocateBuffer(offsets[num_rows]));
+ char *str = reinterpret_cast<char *>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for(int32_t i = 0; i < offsets[num_rows]; i++)
+ str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, { nullptr, std::move(offset_buff), std::move(str_buff) });
+ return std::move(ad);
+ }
+
+ void AppendNumber(char *&out, int num_digits, int32_t x)
+ {
+ out += (num_digits - 1);
+ while(x > 0)
+ {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+ }
+
+ void GeneratePhoneNumber(
+ char *out,
+ random::pcg32_fast &rng,
+ int32_t country)
+ {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+ }
+
+ static constexpr uint32_t kStartDate = 8035; // January 1, 1992 is 8035 days after January 1, 1970
+ static constexpr uint32_t kCurrentDate = 9298; // June 17, 1995 is 9298 days after January 1, 1970
+ static constexpr uint32_t kEndDate = 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+ using GenerateColumnFn = std::function<Status(size_t)>;
+ class PartAndPartSupplierGenerator
+ {
+ public:
+ Status Init(
+ size_t num_threads,
+ int64_t batch_size,
+ float scale_factor)
+ {
+ if(!inited_)
+ {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for(ThreadLocalData &tld : thread_local_data_)
+ {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const
+ {
+ return part_batches_generated_.load();
+ }
+
+ int64_t partsupp_batches_generated() const
+ {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch()
+ {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if(!part_output_queue_.empty())
+ {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ }
+ else if(part_rows_generated_ == part_rows_to_generate_)
+ {
+ return util::nullopt;
+ }
+ else
+ {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate = std::min(
+ batch_size_,
+ part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for(int col : part_cols_)
+ RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for(int col : partsupp_cols_)
+ RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for(size_t i = 0; i < part_cols_.size(); i++)
+ {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if(!partsupp_cols_.empty())
+ {
+ std::vector<ExecBatch> partsupp_results;
+ for(size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++)
+ {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for(size_t icol = 0; icol < partsupp_cols_.size(); icol++)
+ {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for(ExecBatch &eb : partsupp_results)
+ {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch()
+ {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ if(!partsupp_output_queue_.empty())
+ {
+ ExecBatch result = std::move(partsupp_output_queue_.front());
+ partsupp_output_queue_.pop();
+ return std::move(result);
+ }
+ }
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if(part_rows_generated_ == part_rows_to_generate_)
+ {
+ return util::nullopt;
+ }
+ else
+ {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate = std::min(
+ batch_size_,
+ part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for(int col : part_cols_)
+ RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for(int col : partsupp_cols_)
+ RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+ if(!part_cols_.empty())
+ {
+ std::vector<Datum> part_result(part_cols_.size());
+ for(size_t i = 0; i < part_cols_.size(); i++)
+ {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch part_batch, ExecBatch::Make(std::move(part_result)));
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ part_output_queue_.emplace(std::move(part_batch));
+ }
+ }
+ std::vector<ExecBatch> partsupp_results;
+ for(size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++)
+ {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for(size_t icol = 0; icol < partsupp_cols_.size(); icol++)
+ {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ for(size_t i = 1; i < partsupp_results.size(); i++)
+ partsupp_output_queue_.emplace(std::move(partsupp_results[i]));
+ }
+ return std::move(partsupp_results[0]);
+ }
+
+ private:
+#define FOR_EACH_PART_COLUMN(F) \
+ F(P_PARTKEY) \
+ F(P_NAME) \
+ F(P_MFGR) \
+ F(P_BRAND) \
+ F(P_TYPE) \
+ F(P_SIZE) \
+ F(P_CONTAINER) \
+ F(P_RETAILPRICE) \
+ F(P_COMMENT)
+
+#define FOR_EACH_PARTSUPP_COLUMN(F) \
+ F(PS_PARTKEY) \
+ F(PS_SUPPKEY) \
+ F(PS_AVAILQTY) \
+ F(PS_SUPPLYCOST) \
+ F(PS_COMMENT) \
+
+#define MAKE_ENUM(col) col,
+ struct PART
+ {
+ enum
+ {
+ FOR_EACH_PART_COLUMN(MAKE_ENUM)
+ kNumCols,
+ };
+ };
+ struct PARTSUPP
+ {
+ enum
+ {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_ENUM)
+ kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) \
+ { #col, PART::col },
+ const std::unordered_map<std::string, int> kPartNameMap =
+ {
+ FOR_EACH_PART_COLUMN(MAKE_STRING_MAP)
+ };
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) \
+ { #col, PARTSUPP::col },
+ const std::unordered_map<std::string, int> kPartsuppNameMap =
+ {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_STRING_MAP)
+ };
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kPartGenerators =
+ {
+ FOR_EACH_PART_COLUMN(MAKE_FN_ARRAY)
+ };
+ std::vector<GenerateColumnFn> kPartsuppGenerators =
+ {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_FN_ARRAY)
+ };
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kPartTypes =
+ {
+ int32(),
+ utf8(),
+ fixed_size_binary(25),
+ fixed_size_binary(10),
+ utf8(),
+ int32(),
+ fixed_size_binary(10),
+ decimal(12, 2),
+ utf8(),
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kPartsuppTypes =
+ {
+ int32(),
+ int32(),
+ int32(),
+ decimal(12, 2),
+ utf8(),
+ };
+
+ Status AllocatePartBatch(size_t thread_index, int column)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.part[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff, AllocateBuffer(tld.part_to_generate * byte_width));
+ ArrayData ad(kPartTypes[column], tld.part_to_generate, { nullptr, std::move(buff) });
+ tld.part[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status P_PARTKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_PARTKEY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_PARTKEY));
+ int32_t *p_partkey = reinterpret_cast<int32_t *>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.part_to_generate; i++)
+ {
+ p_partkey[i] = static_cast<int32_t>(tld.partkey_start + i + 1);
+ ARROW_DCHECK(1 <= p_partkey[i] && p_partkey[i] <= part_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_NAME(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_NAME].kind() == Datum::NONE)
+ {
+ std::uniform_int_distribution<int32_t> dist(0, static_cast<uint8_t>(kNumNameParts - 1));
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff, AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ size_t string_length = 0;
+ for(int ipart = 0; ipart < 5; ipart++)
+ {
+ uint8_t name_part_index = static_cast<uint8_t>(dist(tld.rng));
+ tld.string_indices[irow * 5 + ipart] = name_part_index;
+ string_length += std::strlen(NameParts[name_part_index]);
+ }
+ // Add 4 because there is a space between each word (i.e. four spaces)
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length + 4);
+ }
+ // Add an extra byte for the space after in the very last string.
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer, AllocateBuffer(offsets[tld.part_to_generate] + 1));
+ char *strings = reinterpret_cast<char *>(string_buffer->mutable_data());
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ char *row = strings + offsets[irow];
+ for(int ipart = 0; ipart < 5; ipart++)
+ {
+ uint8_t name_part_index = tld.string_indices[irow * 5 + ipart];
+ const char *part = NameParts[name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ *row++ = ' ';
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_NAME], tld.part_to_generate, { nullptr, std::move(offset_buff), std::move(string_buffer) });
+ Datum datum(ad);
+ tld.part[PART::P_NAME] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_MFGR(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_MFGR].kind() == Datum::NONE)
+ {
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char *manufacturer = "Manufacturer#";
+ const size_t manufacturer_length = std::strlen(manufacturer);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_MFGR));
+ char *p_mfgr = reinterpret_cast<char *>(tld.part[PART::P_MFGR].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ std::strncpy(p_mfgr + byte_width * irow, manufacturer, byte_width);
+ char mfgr_id = '0' + dist(tld.rng);
+ *(p_mfgr + byte_width * irow + manufacturer_length) = mfgr_id;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_BRAND(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_BRAND].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(P_MFGR(thread_index));
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char *brand = "Brand#";
+ const size_t brand_length = std::strlen(brand);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_BRAND));
+ const char *p_mfgr = reinterpret_cast<const char *>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->data());
+ char *p_brand = reinterpret_cast<char *>(
+ tld.part[PART::P_BRAND].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_BRAND]);
+ int32_t mfgr_byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ const size_t mfgr_id_offset = std::strlen("Manufacturer#");
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ char *row = p_brand + byte_width * irow;
+ char mfgr_id = *(p_mfgr + irow * mfgr_byte_width + mfgr_id_offset);
+ char brand_id = '0' + dist(tld.rng);
+ std::strncpy(row, brand, byte_width);
+ *(row + brand_length) = mfgr_id;
+ *(row + brand_length + 1) = brand_id;
+ irow += 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_TYPE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_TYPE].kind() == Datum::NONE)
+ {
+ using D = std::uniform_int_distribution<uint32_t>;
+ D dists[] =
+ {
+ D{ 0, static_cast<uint8_t>(kNumTypes_1 - 1) },
+ D{ 0, static_cast<uint8_t>(kNumTypes_2 - 1) },
+ D{ 0, static_cast<uint8_t>(kNumTypes_3 - 1) },
+ };
+
+ const char **types[] = { Types_1, Types_2, Types_3 };
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff, AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ size_t string_length = 0;
+ for(int ipart = 0; ipart < 3; ipart++)
+ {
+ uint8_t name_part_index = static_cast<uint8_t>(dists[ipart](tld.rng));
+ tld.string_indices[irow * 3 + ipart] = name_part_index;
+ string_length += std::strlen(types[ipart][name_part_index]);
+ }
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer, AllocateBuffer(offsets[tld.part_to_generate]));
+ char *strings = reinterpret_cast<char *>(string_buffer->mutable_data());
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ char *row = strings + offsets[irow];
+ for(int ipart = 0; ipart < 3; ipart++)
+ {
+ uint8_t name_part_index = tld.string_indices[irow * 3 + ipart];
+ const char *part = types[ipart][name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_TYPE], tld.part_to_generate, { nullptr, std::move(offset_buff), std::move(string_buffer) });
+ Datum datum(ad);
+ tld.part[PART::P_TYPE] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_SIZE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_SIZE].kind() == Datum::NONE)
+ {
+ std::uniform_int_distribution<int32_t> dist(1, 50);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_SIZE));
+ int32_t *p_size = reinterpret_cast<int32_t *>(
+ tld.part[PART::P_SIZE].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.part_to_generate; i++)
+ p_size[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status P_CONTAINER(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_CONTAINER].kind() == Datum::NONE)
+ {
+ std::uniform_int_distribution<int> dist1(0, static_cast<uint8_t>(kNumContainers_1 - 1));
+ std::uniform_int_distribution<int> dist2(0, static_cast<uint8_t>(kNumContainers_2 - 1));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_CONTAINER));
+ char *p_container = reinterpret_cast<char *>(
+ tld.part[PART::P_CONTAINER].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_CONTAINER]);
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ int container1_idx = dist1(tld.rng);
+ int container2_idx = dist2(tld.rng);
+ const char *container1 = Containers_1[container1_idx];
+ const char *container2 = Containers_2[container2_idx];
+ size_t container1_length = std::strlen(container1);
+ size_t container2_length = std::strlen(container2);
+
+ char *row = p_container + byte_width * irow;
+ std::strncpy(row, container1, byte_width);
+ std::memcpy(row + container1_length, container2, container2_length);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_RETAILPRICE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_RETAILPRICE].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_RETAILPRICE));
+ const int32_t *p_partkey = reinterpret_cast<const int32_t *>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+ Decimal128 *p_retailprice = reinterpret_cast<Decimal128 *>(
+ tld.part[PART::P_RETAILPRICE].array()->buffers[1]->mutable_data());
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ int32_t partkey = p_partkey[irow];
+ int64_t retail_price = (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ p_retailprice[irow] = { retail_price };
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_COMMENT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_COMMENT].kind() == Datum::NONE)
+ {
+ ARROW_ASSIGN_OR_RAISE(tld.part[PART::P_COMMENT], g_text.GenerateComments(tld.part_to_generate, 5, 22, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ int64_t PartsuppBatchesToGenerate(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t num_batches = (ps_to_generate + batch_size_ - 1) / batch_size_;
+ return num_batches;
+ }
+
+ Status InitPartsupp(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ tld.generated_partsupp.reset();
+ int64_t num_batches = PartsuppBatchesToGenerate(thread_index);
+ tld.partsupp.resize(num_batches);
+ for(std::vector<Datum> &batch : tld.partsupp)
+ {
+ batch.resize(PARTSUPP::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocatePartSuppBatch(size_t thread_index, size_t ibatch, int column)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartsuppTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff, AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kPartsuppTypes[column], batch_size_, { nullptr, std::move(buff) });
+ tld.partsupp[ibatch][column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status PS_PARTKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_partsupp[PARTSUPP::PS_PARTKEY])
+ {
+ tld.generated_partsupp[PARTSUPP::PS_PARTKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t *p_partkey = reinterpret_cast<const int32_t *>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for(int64_t irow = 0; irow < ps_to_generate; ibatch++)
+ {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_PARTKEY));
+ int32_t *ps_partkey = reinterpret_cast<int32_t *>(
+ tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY].array()->buffers[1]->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; ipartsupp < kPartSuppRowsPerPart && irun < next_run; ipartsupp++, irun++)
+ ps_partkey[batch_offset++] = p_partkey[ipart];
+
+ if(ipartsupp == kPartSuppRowsPerPart)
+ {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_partsupp[PARTSUPP::PS_SUPPKEY])
+ {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t *p_partkey = reinterpret_cast<const int32_t *>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for(int64_t irow = 0; irow < ps_to_generate; ibatch++)
+ {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPKEY));
+ int32_t *ps_suppkey = reinterpret_cast<int32_t *>(
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY].array()->buffers[1]->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; ipartsupp < kPartSuppRowsPerPart && irun < next_run; ipartsupp++, irun++)
+ {
+ int32_t supplier = static_cast<int32_t>(ipartsupp);
+ int32_t partkey = p_partkey[ipart];
+ ps_suppkey[batch_offset++] = (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ if(ipartsupp == kPartSuppRowsPerPart)
+ {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_AVAILQTY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_partsupp[PARTSUPP::PS_AVAILQTY])
+ {
+ tld.generated_partsupp[PARTSUPP::PS_AVAILQTY] = true;
+ std::uniform_int_distribution<int32_t> dist(1, 9999);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for(int64_t irow = 0; irow < ps_to_generate; ibatch++)
+ {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_AVAILQTY));
+ int32_t *ps_availqty = reinterpret_cast<int32_t *>(
+ tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY].array()->buffers[1]->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for(int64_t irun = 0; irun < next_run; irun++)
+ ps_availqty[irun] = dist(tld.rng);
+
+ tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPLYCOST(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST])
+ {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST] = true;
+ std::uniform_int_distribution<int64_t> dist(100, 100000);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for(int64_t irow = 0; irow < ps_to_generate; ibatch++)
+ {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPLYCOST));
+ Decimal128 *ps_supplycost = reinterpret_cast<Decimal128 *>(
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST].array()->buffers[1]->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for(int64_t irun = 0; irun < next_run; irun++)
+ ps_supplycost[irun] = { dist(tld.rng) };
+
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_COMMENT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PARTSUPP::PS_COMMENT].kind() == Datum::NONE)
+ {
+ int64_t irow = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for(size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++)
+ {
+ int64_t num_rows = std::min(batch_size_, ps_to_generate - irow);
+ ARROW_ASSIGN_OR_RAISE(
+ tld.partsupp[ibatch][PARTSUPP::PS_COMMENT], g_text.GenerateComments(num_rows, 49, 198, tld.rng));
+ irow += num_rows;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData
+ {
+ std::vector<Datum> part;
+ std::vector<int8_t> string_indices;
+ int64_t part_to_generate{0};
+ int64_t partkey_start{0};
+
+ std::vector<std::vector<Datum>> partsupp;
+ std::bitset<PARTSUPP::kNumCols> generated_partsupp;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex part_output_queue_mutex_;
+ std::mutex partsupp_output_queue_mutex_;
+ std::queue<ExecBatch> part_output_queue_;
+ std::queue<ExecBatch> partsupp_output_queue_;
+ int64_t batch_size_{0};
+ float scale_factor_{0};
+ int64_t part_rows_to_generate_{0};
+ int64_t part_rows_generated_{0};
+ std::vector<int> part_cols_;
+ std::vector<int> partsupp_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<int64_t> part_batches_generated_ = { 0 };
+ std::atomic<int64_t> partsupp_batches_generated_ = { 0 };
+ static constexpr int64_t kPartSuppRowsPerPart = 4;
+ };
+
+ class OrdersAndLineItemGenerator
+ {
+ public:
+ Status Init(
+ size_t num_threads,
+ int64_t batch_size,
+ float scale_factor)
+ {
+ if(!inited_)
+ {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for(ThreadLocalData &tld : thread_local_data_)
+ {
+ tld.items_per_order.resize(batch_size_);
+ tld.rng.seed(seq);
+ }
+ orders_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000 * 10);
+ }
+ return Status::OK();
+ }
+
+ int64_t orders_batches_generated() const
+ {
+ return orders_batches_generated_.load();
+ }
+
+ int64_t lineitem_batches_generated() const
+ {
+ return lineitem_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetOrdersOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, kOrdersTypes, kOrdersNameMap, orders_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetLineItemOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, kLineitemTypes, kLineitemNameMap, lineitem_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextOrdersBatch()
+ {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if(!orders_output_queue_.empty())
+ {
+ ExecBatch batch = std::move(orders_output_queue_.front());
+ orders_output_queue_.pop();
+ return std::move(batch);
+ }
+ else if(orders_rows_generated_ == orders_rows_to_generate_)
+ {
+ return util::nullopt;
+ }
+ else
+ {
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate = std::min(
+ batch_size_,
+ orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.first_batch_offset = 0;
+ tld.generated_lineitem.reset();
+
+ for(int col : orders_cols_)
+ RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for(int col : lineitem_cols_)
+ RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for(size_t i = 0; i < orders_cols_.size(); i++)
+ {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ if(!lineitem_cols_.empty())
+ {
+ std::vector<ExecBatch> lineitem_results;
+ for(size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++)
+ {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for(size_t icol = 0; icol < lineitem_cols_.size(); icol++)
+ {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(lineitem_output_queue_mutex_);
+ for(ExecBatch &eb : lineitem_results)
+ {
+ lineitem_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(orders_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextLineItemBatch()
+ {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ ExecBatch queued;
+ bool from_queue = false;
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ if(!lineitem_output_queue_.empty())
+ {
+ queued = std::move(lineitem_output_queue_.front());
+ lineitem_output_queue_.pop();
+ from_queue = true;
+ }
+ }
+ tld.first_batch_offset = 0;
+ if(from_queue)
+ {
+ ARROW_DCHECK(queued.length <= batch_size_);
+ tld.first_batch_offset = queued.length;
+ if(queued.length == batch_size_)
+ return std::move(queued);
+ }
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if(orders_rows_generated_ == orders_rows_to_generate_)
+ {
+ if(from_queue)
+ return std::move(queued);
+ return util::nullopt;
+ }
+
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate = std::min(
+ batch_size_,
+ orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1ll);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.generated_lineitem.reset();
+ if(from_queue)
+ {
+ lineitem_batches_generated_.fetch_sub(1);
+ for(size_t i = 0; i < lineitem_cols_.size(); i++)
+ if(tld.lineitem[0][lineitem_cols_[i]].kind() == Datum::NONE)
+ tld.lineitem[0][lineitem_cols_[i]] = std::move(queued[i]);
+ }
+
+ for(int col : orders_cols_)
+ RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for(int col : lineitem_cols_)
+ RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ if(!orders_cols_.empty())
+ {
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for(size_t i = 0; i < orders_cols_.size(); i++)
+ {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch orders_batch, ExecBatch::Make(std::move(orders_result)));
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ orders_output_queue_.emplace(std::move(orders_batch));
+ }
+ }
+ std::vector<ExecBatch> lineitem_results;
+ for(size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++)
+ {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for(size_t icol = 0; icol < lineitem_cols_.size(); icol++)
+ {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ lineitem_batches_generated_.fetch_add(static_cast<int64_t>(lineitem_results.size()));
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ for(size_t i = 1; i < lineitem_results.size(); i++)
+ lineitem_output_queue_.emplace(std::move(lineitem_results[i]));
+ }
+ return std::move(lineitem_results[0]);
+ }
+
+ private:
+#define FOR_EACH_ORDERS_COLUMN(F) \
+ F(O_ORDERKEY) \
+ F(O_CUSTKEY) \
+ F(O_ORDERSTATUS) \
+ F(O_TOTALPRICE) \
+ F(O_ORDERDATE) \
+ F(O_ORDERPRIORITY) \
+ F(O_CLERK) \
+ F(O_SHIPPRIORITY) \
+ F(O_COMMENT)
+
+#define FOR_EACH_LINEITEM_COLUMN(F) \
+ F(L_ORDERKEY) \
+ F(L_PARTKEY) \
+ F(L_SUPPKEY) \
+ F(L_LINENUMBER) \
+ F(L_QUANTITY) \
+ F(L_EXTENDEDPRICE) \
+ F(L_DISCOUNT) \
+ F(L_TAX) \
+ F(L_RETURNFLAG) \
+ F(L_LINESTATUS) \
+ F(L_SHIPDATE) \
+ F(L_COMMITDATE) \
+ F(L_RECEIPTDATE) \
+ F(L_SHIPINSTRUCT) \
+ F(L_SHIPMODE) \
+ F(L_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct ORDERS
+ {
+ enum
+ {
+ FOR_EACH_ORDERS_COLUMN(MAKE_ENUM)
+ kNumCols,
+ };
+ };
+ struct LINEITEM
+ {
+ enum
+ {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_ENUM)
+ kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) \
+ { #col, ORDERS::col },
+ const std::unordered_map<std::string, int> kOrdersNameMap =
+ {
+ FOR_EACH_ORDERS_COLUMN(MAKE_STRING_MAP)
+ };
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) \
+ { #col, LINEITEM::col },
+ const std::unordered_map<std::string, int> kLineitemNameMap =
+ {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_STRING_MAP)
+ };
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ const std::vector<GenerateColumnFn> kOrdersGenerators =
+ {
+ FOR_EACH_ORDERS_COLUMN(MAKE_FN_ARRAY)
+ };
+ const std::vector<GenerateColumnFn> kLineitemGenerators =
+ {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_FN_ARRAY)
+ };
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kOrdersTypes =
+ {
+ int32(),
+ int32(),
+ fixed_size_binary(1),
+ decimal(12, 2),
+ date32(),
+ fixed_size_binary(15),
+ fixed_size_binary(15),
+ int32(),
+ utf8()
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kLineitemTypes =
+ {
+ int32(),
+ int32(),
+ int32(),
+ int32(),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ fixed_size_binary(1),
+ fixed_size_binary(1),
+ date32(),
+ date32(),
+ date32(),
+ fixed_size_binary(25),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status AllocateOrdersBatch(size_t thread_index, int column)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.orders[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff, AllocateBuffer(tld.orders_to_generate * byte_width));
+ ArrayData ad(kOrdersTypes[column], tld.orders_to_generate, { nullptr, std::move(buff) });
+ tld.orders[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status O_ORDERKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_ORDERKEY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERKEY));
+ int32_t *o_orderkey = reinterpret_cast<int32_t *>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.orders_to_generate; i++)
+ {
+ int32_t orderkey_index = static_cast<int32_t>(tld.orderkey_start + i);
+ int32_t index_of_run = orderkey_index / 8;
+ int32_t index_in_run = orderkey_index % 8;
+ o_orderkey[i] = (index_of_run * 32 + index_in_run + 1);
+ ARROW_DCHECK(1 <= o_orderkey[i] && o_orderkey[i] <= 4 * orders_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CUSTKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_CUSTKEY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CUSTKEY));
+
+ // Spec says it must be a random number between 1 and SF*150000 that is not
+ // divisible by 3. Rather than repeatedly generating numbers until we get to
+ // a non-divisible-by-3 number, we just generate a number between
+ // 0 and SF * 50000 - 1, multiply by 3, and then add either 1 or 2.
+ int32_t sf_50k = static_cast<int32_t>(scale_factor_ * 50000);
+ std::uniform_int_distribution<int32_t> base_dist(0, sf_50k - 1);
+ std::uniform_int_distribution<int32_t> offset_dist(1, 2);
+ int32_t *o_custkey = reinterpret_cast<int32_t *>(
+ tld.orders[ORDERS::O_CUSTKEY].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.orders_to_generate; i++)
+ o_custkey[i] = 3 * base_dist(tld.rng) + offset_dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERSTATUS(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_ORDERSTATUS].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(L_LINESTATUS(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERSTATUS));
+
+ char *o_orderstatus = reinterpret_cast<char *>(
+ tld.orders[ORDERS::O_ORDERSTATUS].array()->buffers[1]->mutable_data());
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ bool all_f = true;
+ bool all_o = true;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ const char *l_linestatus = reinterpret_cast<const char *>(
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->buffers[1]->data());
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++, batch_offset++)
+ {
+ all_f &= l_linestatus[batch_offset] == 'F';
+ all_o &= l_linestatus[batch_offset] == 'O';
+ }
+ if(iline == tld.items_per_order[iorder])
+ {
+ iline = 0;
+ ARROW_DCHECK(!(all_f && all_o));
+ if(all_f)
+ o_orderstatus[iorder] = 'F';
+ else if(all_o)
+ o_orderstatus[iorder] = 'O';
+ else
+ o_orderstatus[iorder] = 'P';
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_TOTALPRICE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_TOTALPRICE].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(L_EXTENDEDPRICE(thread_index));
+ RETURN_NOT_OK(L_TAX(thread_index));
+ RETURN_NOT_OK(L_DISCOUNT(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_TOTALPRICE));
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ int64_t sum = 0;
+ Decimal128 *o_totalprice = reinterpret_cast<Decimal128 *>(
+ tld.orders[ORDERS::O_TOTALPRICE].array()->buffers[1]->mutable_data());
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ const Decimal128 *l_extendedprice = reinterpret_cast<const Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->buffers[1]->data());
+ const Decimal128 *l_tax = reinterpret_cast<const Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->data());
+ const Decimal128 *l_discount = reinterpret_cast<const Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->buffers[1]->data());
+
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++, batch_offset++)
+ {
+ int64_t eprice = static_cast<int64_t>(l_extendedprice[batch_offset]);
+ int64_t tax = static_cast<int64_t>(l_tax[batch_offset]);
+ int64_t discount = static_cast<int64_t>(l_discount[batch_offset]);
+ sum += (eprice * (100 + tax) * (100 - discount));
+ }
+ if(iline == tld.items_per_order[iorder])
+ {
+ sum /= 100 * 100;
+ o_totalprice[iorder] = { sum };
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERDATE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_ORDERDATE].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERDATE));
+
+ std::uniform_int_distribution<uint32_t> dist(kStartDate, kEndDate - 151);
+ uint32_t *o_orderdate = reinterpret_cast<uint32_t *>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.orders_to_generate; i++)
+ o_orderdate[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERPRIORITY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_ORDERPRIORITY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERPRIORITY));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_ORDERPRIORITY]);
+ std::uniform_int_distribution<int32_t> dist(0, kNumPriorities - 1);
+ char *o_orderpriority = reinterpret_cast<char *>(
+ tld.orders[ORDERS::O_ORDERPRIORITY].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.orders_to_generate; i++)
+ {
+ const char *str = Priorities[dist(tld.rng)];
+ std::strncpy(o_orderpriority + i * byte_width, str, byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CLERK(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_CLERK].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CLERK));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_CLERK]);
+ int64_t max_clerk_id = static_cast<int64_t>(scale_factor_ * 1000);
+ std::uniform_int_distribution<int64_t> dist(1, max_clerk_id);
+ char *o_clerk = reinterpret_cast<char *>(
+ tld.orders[ORDERS::O_CLERK].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.orders_to_generate; i++)
+ {
+ const char *clerk = "Clerk#";
+ const size_t clerk_length = std::strlen(clerk);
+ int64_t clerk_number = dist(tld.rng);
+ char *output = o_clerk + i * byte_width;
+ std::strncpy(output, clerk, byte_width);
+ AppendNumberPaddedToNineDigits(output + clerk_length, clerk_number);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_SHIPPRIORITY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_SHIPPRIORITY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_SHIPPRIORITY));
+ int32_t *o_shippriority = reinterpret_cast<int32_t *>(
+ tld.orders[ORDERS::O_SHIPPRIORITY].array()->buffers[1]->mutable_data());
+ std::memset(o_shippriority, 0, tld.orders_to_generate * sizeof(int32_t));
+ }
+ return Status::OK();
+ }
+
+ Status O_COMMENT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_COMMENT].kind() == Datum::NONE)
+ {
+ ARROW_ASSIGN_OR_RAISE(tld.orders[ORDERS::O_COMMENT], g_text.GenerateComments(tld.orders_to_generate, 19, 78, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ Status GenerateRowCounts(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ std::uniform_int_distribution<int> length_dist(1, 7);
+ tld.lineitem_to_generate = 0;
+ tld.items_per_order.clear();
+ for(int64_t i = 0; i < tld.orders_to_generate; i++)
+ {
+ int length = length_dist(tld.rng);
+ tld.items_per_order.push_back(length);
+ tld.lineitem_to_generate += length;
+ }
+ int64_t num_batches = (tld.first_batch_offset + tld.lineitem_to_generate + batch_size_ - 1) / batch_size_;
+ tld.lineitem.resize(num_batches);
+ for(std::vector<Datum> &batch : tld.lineitem)
+ {
+ batch.resize(LINEITEM::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocateLineItemBufferIfNeeded(size_t thread_index, size_t ibatch, int column, size_t &out_batch_offset)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.lineitem[ibatch][column].kind() == Datum::NONE)
+ {
+ int32_t byte_width = arrow::internal::GetByteWidth(*kLineitemTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff, AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kLineitemTypes[column], batch_size_, { nullptr, std::move(buff) });
+ tld.lineitem[ibatch][column] = std::move(ad);
+ out_batch_offset = 0;
+ }
+ if(ibatch == 0)
+ out_batch_offset = tld.first_batch_offset;
+
+ return Status::OK();
+ }
+
+ Status L_ORDERKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_ORDERKEY])
+ {
+ tld.generated_lineitem[LINEITEM::L_ORDERKEY] = true;
+ RETURN_NOT_OK(O_ORDERKEY(thread_index));
+ const int32_t *o_orderkey = reinterpret_cast<const int32_t *>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_ORDERKEY, batch_offset));
+ int32_t *l_linenumber = reinterpret_cast<int32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_ORDERKEY].array()->buffers[1]->mutable_data());
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_linenumber[batch_offset++] = o_orderkey[iorder];
+ if(iline == tld.items_per_order[iorder])
+ {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_ORDERKEY].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_PARTKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_PARTKEY])
+ {
+ tld.generated_lineitem[LINEITEM::L_PARTKEY] = true;
+
+ size_t ibatch = 0;
+ int32_t max_partkey = static_cast<int32_t>(scale_factor_ * 200000);
+ std::uniform_int_distribution<int32_t> dist(1, max_partkey);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_PARTKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t *l_partkey = reinterpret_cast<int32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_partkey[batch_offset] = dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SUPPKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_SUPPKEY])
+ {
+ tld.generated_lineitem[LINEITEM::L_SUPPKEY] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 3);
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset = 0;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_SUPPKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t *l_suppkey = reinterpret_cast<int32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_SUPPKEY].array()->buffers[1]->mutable_data());
+ const int32_t *l_partkey = reinterpret_cast<const int32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ for(int64_t i = 0; i < next_run; i++)
+ {
+ int32_t supplier = dist(tld.rng);
+ int32_t partkey = l_partkey[batch_offset];
+ // Fun fact: the parentheses for this expression are unbalanced in the TPC-H spec.
+ l_suppkey[batch_offset++] = (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SUPPKEY].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINENUMBER(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_LINENUMBER])
+ {
+ tld.generated_lineitem[LINEITEM::L_LINENUMBER] = true;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_LINENUMBER, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t *l_linenumber = reinterpret_cast<int32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_LINENUMBER].array()->buffers[1]->mutable_data());
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ {
+ l_linenumber[batch_offset++] = (iline + 1);
+ ARROW_DCHECK(1 <= (iline + 1) && (iline + 1) <= 7);
+ }
+ if(iline == tld.items_per_order[iorder])
+ {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINENUMBER].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_QUANTITY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_QUANTITY])
+ {
+ tld.generated_lineitem[LINEITEM::L_QUANTITY] = true;
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int64_t> dist(1, 50);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_QUANTITY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128 *l_quantity = reinterpret_cast<Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < next_run; i++)
+ {
+ // Multiply by 100 because the type is decimal(12, 2), so the decimal goes after two digits
+ int64_t quantity = dist(tld.rng) * 100;
+ l_quantity[batch_offset++] = { quantity };
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_EXTENDEDPRICE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE])
+ {
+ tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+ RETURN_NOT_OK(L_QUANTITY(thread_index));
+ size_t ibatch = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_EXTENDEDPRICE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ const int32_t *l_partkey = reinterpret_cast<const int32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ const Decimal128 *l_quantity = reinterpret_cast<const Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->buffers[1]->data());
+ Decimal128 *l_extendedprice = reinterpret_cast<Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ {
+ int64_t partkey = static_cast<int64_t>(l_partkey[batch_offset]);
+ // Divide by 100 to recover the integer representation (not Decimal).
+ int64_t quantity = static_cast<int64_t>(l_quantity[batch_offset]) / 100;
+
+ // Spec says to divide by 100, but that happens automatically due to this being stored
+ // to two decimal points.
+ int64_t retail_price = (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ int64_t extended_price = retail_price * quantity;
+ l_extendedprice[batch_offset] = { extended_price };
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_DISCOUNT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_DISCOUNT])
+ {
+ tld.generated_lineitem[LINEITEM::L_DISCOUNT] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 10);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_DISCOUNT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128 *l_discount = reinterpret_cast<Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_discount[batch_offset] = { dist(tld.rng) };
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_TAX(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_TAX])
+ {
+ tld.generated_lineitem[LINEITEM::L_TAX] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 8);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_TAX, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128 *l_tax = reinterpret_cast<Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_tax[batch_offset] = { dist(tld.rng) };
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RETURNFLAG(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_RETURNFLAG])
+ {
+ tld.generated_lineitem[LINEITEM::L_RETURNFLAG] = true;
+ RETURN_NOT_OK(L_RECEIPTDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<uint32_t> dist;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_RETURNFLAG, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char *l_returnflag = reinterpret_cast<char *>(
+ tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG].array()->buffers[1]->mutable_data());
+ const uint32_t *l_receiptdate = reinterpret_cast<const uint32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ {
+ if(l_receiptdate[batch_offset] <= kCurrentDate)
+ {
+ uint32_t r = dist(tld.rng);
+ l_returnflag[batch_offset] = (r % 2 == 1) ? 'R' : 'A';
+ }
+ else
+ {
+ l_returnflag[batch_offset] = 'N';
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINESTATUS(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_LINESTATUS])
+ {
+ tld.generated_lineitem[LINEITEM::L_LINESTATUS] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_LINESTATUS, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char *l_linestatus = reinterpret_cast<char *>(
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->buffers[1]->mutable_data());
+ const uint32_t *l_shipdate = reinterpret_cast<const uint32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ {
+ if(l_shipdate[batch_offset] > kCurrentDate)
+ l_linestatus[batch_offset] = 'O';
+ else
+ l_linestatus[batch_offset] = 'F';
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPDATE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_SHIPDATE])
+ {
+ tld.generated_lineitem[LINEITEM::L_SHIPDATE] = true;
+ RETURN_NOT_OK(O_ORDERDATE(thread_index));
+ const int32_t *o_orderdate = reinterpret_cast<const int32_t *>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(1, 121);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_SHIPDATE, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t *l_shipdate = reinterpret_cast<uint32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->buffers[1]->mutable_data());
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_shipdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if(iline == tld.items_per_order[iorder])
+ {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMITDATE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_COMMITDATE])
+ {
+ tld.generated_lineitem[LINEITEM::L_COMMITDATE] = true;
+ const int32_t *o_orderdate = reinterpret_cast<const int32_t *>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(30, 90);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_COMMITDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t *l_commitdate = reinterpret_cast<uint32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_COMMITDATE].array()->buffers[1]->mutable_data());
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_commitdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if(iline == tld.items_per_order[iorder])
+ {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMITDATE].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RECEIPTDATE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_RECEIPTDATE])
+ {
+ tld.generated_lineitem[LINEITEM::L_RECEIPTDATE] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(1, 30);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_RECEIPTDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t *l_receiptdate = reinterpret_cast<uint32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE].array()->buffers[1]->mutable_data());
+ const uint32_t *l_shipdate = reinterpret_cast<const uint32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_receiptdate[batch_offset] = l_shipdate[batch_offset] + dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPINSTRUCT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT])
+ {
+ tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT] = true;
+ int32_t byte_width = arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPINSTRUCT]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumInstructions - 1);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_SHIPINSTRUCT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char *l_shipinstruct = reinterpret_cast<char *>(
+ tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ {
+ const char *str = Instructions[dist(tld.rng)];
+ // Note that we don't have to memset the buffer to 0 because strncpy pads each string
+ // with 0's anyway
+ std::strncpy(l_shipinstruct + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPMODE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_SHIPMODE])
+ {
+ tld.generated_lineitem[LINEITEM::L_SHIPMODE] = true;
+ int32_t byte_width = arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPMODE]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumModes - 1);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_SHIPMODE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char *l_shipmode = reinterpret_cast<char *>(
+ tld.lineitem[ibatch][LINEITEM::L_SHIPMODE].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ {
+ const char *str = Modes[dist(tld.rng)];
+ std::strncpy(l_shipmode + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPMODE].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMENT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_COMMENT])
+ {
+ tld.generated_lineitem[LINEITEM::L_COMMENT] = true;
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ // Comments are kind of sneaky: we always generate the full batch and then just bump the length
+ if(tld.lineitem[ibatch][LINEITEM::L_COMMENT].kind() == Datum::NONE)
+ {
+ ARROW_ASSIGN_OR_RAISE(tld.lineitem[ibatch][LINEITEM::L_COMMENT], g_text.GenerateComments(batch_size_, 10, 43, tld.rng));
+ batch_offset = 0;
+ }
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ batch_offset += next_run;
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMENT].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData
+ {
+ std::vector<Datum> orders;
+ int64_t orders_to_generate;
+ int64_t orderkey_start;
+
+ std::vector<std::vector<Datum>> lineitem;
+ std::vector<int> items_per_order;
+ int64_t lineitem_to_generate;
+ int64_t first_batch_offset;
+ std::bitset<LINEITEM::kNumCols> generated_lineitem;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex orders_output_queue_mutex_;
+ std::mutex lineitem_output_queue_mutex_;
+ std::queue<ExecBatch> orders_output_queue_;
+ std::queue<ExecBatch> lineitem_output_queue_;
+ int64_t batch_size_;
+ float scale_factor_;
+ int64_t orders_rows_to_generate_;
+ int64_t orders_rows_generated_;
+ std::vector<int> orders_cols_;
+ std::vector<int> lineitem_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<size_t> orders_batches_generated_ = { 0 };
+ std::atomic<size_t> lineitem_batches_generated_ = { 0 };
+ };
+
+ class SupplierGenerator : public TpchTableGenerator
+ {
+ public:
+ Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) override
+ {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 10000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_, SetOutputColumns(
+ columns,
+ kTypes,
+ kNameMap,
+ gen_list_));
+
+ random::pcg32_fast rng;
+ std::uniform_int_distribution<int64_t> dist(0, rows_to_generate_ - 1);
+ size_t num_special_rows = static_cast<size_t>(5 * scale_factor_);
+ std::unordered_set<int64_t> good_rows_set;
+ while(good_rows_set.size() < num_special_rows)
+ {
+ int64_t row = dist(rng);
+ good_rows_set.insert(row);
+ }
+ std::unordered_set<int64_t> bad_rows_set;
+ while(bad_rows_set.size() < num_special_rows)
+ {
+ int64_t bad_row;
+ do
+ {
+ bad_row = dist(rng);
+ } while(good_rows_set.find(bad_row) != good_rows_set.end());
+ bad_rows_set.insert(bad_row);
+ }
+ good_rows_.clear();
+ bad_rows_.clear();
+ good_rows_.insert(good_rows_.end(), good_rows_set.begin(), good_rows_set.end());
+ bad_rows_.insert(bad_rows_.end(), bad_rows_set.begin(), bad_rows_set.end());
+ std::sort(good_rows_.begin(), good_rows_.end());
+ std::sort(bad_rows_.begin(), bad_rows_.end());
+ return Status::OK();
+ }
+
+ Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override
+ {
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for(ThreadLocalData &tld : thread_local_data_)
+ tld.rng.seed(seq);
+
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for(size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override
+ {
+ return schema_;
+ }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(S_SUPPKEY) \
+ F(S_NAME) \
+ F(S_ADDRESS) \
+ F(S_NATIONKEY) \
+ F(S_PHONE) \
+ F(S_ACCTBAL) \
+ F(S_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct SUPPLIER
+ {
+ enum
+ {
+ FOR_EACH_COLUMN(MAKE_ENUM)
+ kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) \
+ { #col, SUPPLIER::col },
+ const std::unordered_map<std::string, int> kNameMap =
+ {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)
+ };
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kGenerators =
+ {
+ FOR_EACH_COLUMN(MAKE_FN_ARRAY)
+ };
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> kTypes =
+ {
+ int32(),
+ fixed_size_binary(25),
+ utf8(),
+ int32(),
+ fixed_size_binary(15),
+ decimal(12, 2),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index)
+ {
+ if(done_.load())
+ return Status::OK();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ tld.suppkey_start = rows_generated_.fetch_add(batch_size_);
+ if(tld.suppkey_start >= rows_to_generate_)
+ return Status::OK();
+
+ tld.to_generate = std::min(batch_size_,
+ rows_to_generate_ - tld.suppkey_start);
+
+ tld.batch.resize(SUPPLIER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for(int col : gen_list_)
+ RETURN_NOT_OK(kGenerators[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for(size_t i = 0; i < gen_list_.size(); i++)
+ {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_outputted_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_outputted_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if(is_last_batch)
+ {
+ done_.store(true);
+ finished_callback_(batches_outputted_.load());
Review comment:
We should set `done_` with `compare_exchange_strong` like we do elsewhere and guard the call to `finished_callback_`
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3850 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ const char *NameParts[] =
+ {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque", "black", "blanched", "blue",
+ "blush", "brown", "burlywood", "burnished", "chartreuse", "chiffon", "chocolate", "coral",
+ "cornflower", "cornsilk", "cream", "cyan", "dark", "deep", "dim", "dodger", "drab", "firebrick",
+ "floral", "forest", "frosted", "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender", "lawn", "lemon", "light", "lime", "linen",
+ "magenta", "maroon", "medium", "metallic", "midnight", "mint", "misty", "moccasin", "navajo",
+ "navy", "olive", "orange", "orchid", "pale", "papaya", "peach", "peru", "pink", "plum", "powder",
+ "puff", "purple", "red", "rose", "rosy", "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel", "tan", "thistle", "tomato", "turquoise", "violet",
+ "wheat", "white", "yellow",
+ };
+ static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+ const char *Types_1[] =
+ {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+ };
+ static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+ const char *Types_2[] =
+ {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+ };
+ static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+ const char *Types_3[] =
+ {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+ };
+ static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+ const char *Containers_1[] =
+ {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+ };
+ static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+ const char *Containers_2[] =
+ {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+ };
+ static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+ const char *Segments[] =
+ {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+ };
+ static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+ const char *Priorities[] =
+ {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+ };
+ static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+ const char *Instructions[] =
+ {
+ "DELIVER IN PERSON", "COLLECT COD", "NONE", "TAKE BACK RETURN",
+ };
+ static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+ const char *Modes[] =
+ {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+ };
+ static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+ const char *Nouns[] =
+ {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ", "dependencies ", "excuses ",
+ "platelets ", "asymptotes ", "courts ", "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ", "forges ", "braids ", "hockey players ", "frays ",
+ "warhorses ", "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ", "waters ", "orbits ", "gifts ", "sheaves ",
+ "depths ", "sentiments ", "decoys ", "realms ", "pains ", "grouches ", "escapades ",
+ };
+ static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+ const char *Verbs[] =
+ {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ", "boost ", "affix ", "detect ", "integrate ",
+ "maintain ", "nod ", "was ", "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ", "poach ", "serve ", "run ", "dazzle ",
+ "snooze ", "doze ", "unwind ", "kindle ", "play ", "hang ", "believe ", "doubt ",
+ };
+ static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+ const char *Adjectives[] =
+ {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ", "quiet ", "ruthless ", "thin ",
+ "close ", "dogged ", "daring ", "brave ", "stealthy ", "permanent ", "enticing ", "idle ", "busy ",
+ "regular ", "final ", "ironic ", "even ", "bold ", "silent ",
+ };
+ static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+ const char *Adverbs[] =
+ {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ", "blithely ", "quickly ", "fluffily ",
+ "slowly ", "quietly ", "ruthlessly ", "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ", "ironically ", "evenly ", "boldly ",
+ "silently ",
+ };
+ static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+ const char *Prepositions[] =
+ {
+ "about ", "above ", "according to ", "across ", "after ", "against ", "along ", "alongside of ", "among ",
+ "around ", "at ", "atop ", "before ", "behind ", "beneath ", "beside ", "besides ", "between ", "beyond ",
+ "beyond ", "by ", "despite ", "during ", "except ", "for ", "from ", "in place of ", "inside ", "instead of ",
+ "into ", "near ", "of ", "on ", "outside ", "over ", "past ", "since ", "through ", "throughout ", "to ",
+ "toward ", "under ", "until ", "up ", "upon ", "without ", "with ", "within ",
+ };
+ static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+ const char *Auxiliaries[] =
+ {
+ "do ", "may ", "might ", "shall ", "will ", "would ", "can ", "could ", "should ", "ought to ", "must ",
+ "will have to ", "shall have to ", "could have to ", "should have to ", "must have to ", "need to ", "try to ",
+ };
+ static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+ const char *Terminators[] =
+ {
+ ".", ";", ":", "?", "!", "--",
+ };
+ static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+
+ // The spec says to generate a 300 MB string according to a grammar. This is a
+ // concurrent implementation of the generator. Each thread generates the text in
+ // (up to) 8KB chunks of text. The generator maintains a cursor into the
+ // 300 MB buffer. After generating the chunk, the cursor is incremented
+ // to reserve space, and the chunk is memcpy-d in.
+ // This text is used to generate the COMMENT columns. To generate a comment, the spec
+ // says to pick a random length and a random offset into the 300 MB buffer (it does
+ // not specify it should be word/sentence aligned), and that slice of text becomes
+ // the comment.
+ class TpchPseudotext
+ {
+ public:
+ Status EnsureInitialized(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ static TpchPseudotext g_text;
+
+ Status TpchPseudotext::EnsureInitialized(random::pcg32_fast &rng)
+ {
+ if(done_.load())
+ return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(!text_)
+ {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char *out = reinterpret_cast<char *>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while(!done_.load())
+ {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while(GenerateSentence(try_offset, rng, temp_buff))
+ known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(done_.load())
+ return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if(last_one)
+ done_.store(true);
+ }
+ return Status::OK();
+ }
+
+ Result<Datum> TpchPseudotext::GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng)
+ {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer, AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for(size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer, AllocateBuffer(offsets[num_comments]));
+ char *comments = reinterpret_cast<char *>(comment_buffer->mutable_data());
+ for(size_t i = 0; i < num_comments; i++)
+ {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments, { nullptr, std::move(offset_buffer), std::move(comment_buffer) });
+ return std::move(ad);
+ }
+
+ bool TpchPseudotext::GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices)
+ {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char *word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if(offset + length > kChunkSize)
+ return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+ }
+
+ bool TpchPseudotext::GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+ }
+
+ bool TpchPseudotext::GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+ }
+
+ bool TpchPseudotext::GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+ }
+
+ bool TpchPseudotext::GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+ }
+
+ bool TpchPseudotext::GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+ }
+
+ bool TpchPseudotext::GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+ }
+
+ bool TpchPseudotext::GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if(result)
+ std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+ }
+
+ bool TpchPseudotext::GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char *comma_space = ", ";
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ bool TpchPseudotext::GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ const char *the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+ }
+
+ bool TpchPseudotext::GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort()
+ {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string> &columns,
+ const std::vector<std::shared_ptr<DataType>> &types,
+ const std::unordered_map<std::string, int> &name_map,
+ std::vector<int> &gen_list)
+ {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if(columns.empty())
+ {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for(auto pair : name_map)
+ {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ }
+ else
+ {
+ for(const std::string &col : columns)
+ {
+ auto entry = name_map.find(col);
+ if(entry == name_map.end())
+ return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+ }
+
+ Result<Datum> RandomVString(
+ random::pcg32_fast &rng,
+ int64_t num_rows,
+ int32_t min_length,
+ int32_t max_length)
+ {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff, AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for(int64_t i = 1; i <= num_rows; i++)
+ offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff, AllocateBuffer(offsets[num_rows]));
+ char *str = reinterpret_cast<char *>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for(int32_t i = 0; i < offsets[num_rows]; i++)
+ str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, { nullptr, std::move(offset_buff), std::move(str_buff) });
+ return std::move(ad);
+ }
+
+ void AppendNumber(char *&out, int num_digits, int32_t x)
+ {
+ out += (num_digits - 1);
+ while(x > 0)
+ {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+ }
+
+ void GeneratePhoneNumber(
+ char *out,
+ random::pcg32_fast &rng,
+ int32_t country)
+ {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+ }
+
+ static constexpr uint32_t kStartDate = 8035; // January 1, 1992 is 8035 days after January 1, 1970
+ static constexpr uint32_t kCurrentDate = 9298; // June 17, 1995 is 9298 days after January 1, 1970
+ static constexpr uint32_t kEndDate = 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+ using GenerateColumnFn = std::function<Status(size_t)>;
+ class PartAndPartSupplierGenerator
+ {
+ public:
+ Status Init(
+ size_t num_threads,
+ int64_t batch_size,
+ float scale_factor)
+ {
+ if(!inited_)
+ {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for(ThreadLocalData &tld : thread_local_data_)
+ {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const
+ {
+ return part_batches_generated_.load();
+ }
+
+ int64_t partsupp_batches_generated() const
+ {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch()
+ {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if(!part_output_queue_.empty())
+ {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ }
+ else if(part_rows_generated_ == part_rows_to_generate_)
+ {
+ return util::nullopt;
+ }
+ else
+ {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate = std::min(
+ batch_size_,
+ part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for(int col : part_cols_)
+ RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for(int col : partsupp_cols_)
+ RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for(size_t i = 0; i < part_cols_.size(); i++)
+ {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if(!partsupp_cols_.empty())
+ {
+ std::vector<ExecBatch> partsupp_results;
+ for(size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++)
+ {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for(size_t icol = 0; icol < partsupp_cols_.size(); icol++)
+ {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for(ExecBatch &eb : partsupp_results)
+ {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch()
+ {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ if(!partsupp_output_queue_.empty())
+ {
+ ExecBatch result = std::move(partsupp_output_queue_.front());
+ partsupp_output_queue_.pop();
+ return std::move(result);
+ }
+ }
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if(part_rows_generated_ == part_rows_to_generate_)
+ {
+ return util::nullopt;
+ }
+ else
+ {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate = std::min(
+ batch_size_,
+ part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for(int col : part_cols_)
+ RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for(int col : partsupp_cols_)
+ RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+ if(!part_cols_.empty())
+ {
+ std::vector<Datum> part_result(part_cols_.size());
+ for(size_t i = 0; i < part_cols_.size(); i++)
+ {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch part_batch, ExecBatch::Make(std::move(part_result)));
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ part_output_queue_.emplace(std::move(part_batch));
+ }
+ }
+ std::vector<ExecBatch> partsupp_results;
+ for(size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++)
+ {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for(size_t icol = 0; icol < partsupp_cols_.size(); icol++)
+ {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ for(size_t i = 1; i < partsupp_results.size(); i++)
+ partsupp_output_queue_.emplace(std::move(partsupp_results[i]));
+ }
+ return std::move(partsupp_results[0]);
+ }
+
+ private:
+#define FOR_EACH_PART_COLUMN(F) \
+ F(P_PARTKEY) \
+ F(P_NAME) \
+ F(P_MFGR) \
+ F(P_BRAND) \
+ F(P_TYPE) \
+ F(P_SIZE) \
+ F(P_CONTAINER) \
+ F(P_RETAILPRICE) \
+ F(P_COMMENT)
+
+#define FOR_EACH_PARTSUPP_COLUMN(F) \
+ F(PS_PARTKEY) \
+ F(PS_SUPPKEY) \
+ F(PS_AVAILQTY) \
+ F(PS_SUPPLYCOST) \
+ F(PS_COMMENT) \
+
+#define MAKE_ENUM(col) col,
+ struct PART
+ {
+ enum
+ {
+ FOR_EACH_PART_COLUMN(MAKE_ENUM)
+ kNumCols,
+ };
+ };
+ struct PARTSUPP
+ {
+ enum
+ {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_ENUM)
+ kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) \
+ { #col, PART::col },
+ const std::unordered_map<std::string, int> kPartNameMap =
+ {
+ FOR_EACH_PART_COLUMN(MAKE_STRING_MAP)
+ };
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) \
+ { #col, PARTSUPP::col },
+ const std::unordered_map<std::string, int> kPartsuppNameMap =
+ {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_STRING_MAP)
+ };
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kPartGenerators =
+ {
+ FOR_EACH_PART_COLUMN(MAKE_FN_ARRAY)
+ };
+ std::vector<GenerateColumnFn> kPartsuppGenerators =
+ {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_FN_ARRAY)
+ };
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kPartTypes =
+ {
+ int32(),
+ utf8(),
+ fixed_size_binary(25),
+ fixed_size_binary(10),
+ utf8(),
+ int32(),
+ fixed_size_binary(10),
+ decimal(12, 2),
+ utf8(),
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kPartsuppTypes =
+ {
+ int32(),
+ int32(),
+ int32(),
+ decimal(12, 2),
+ utf8(),
+ };
+
+ Status AllocatePartBatch(size_t thread_index, int column)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.part[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff, AllocateBuffer(tld.part_to_generate * byte_width));
+ ArrayData ad(kPartTypes[column], tld.part_to_generate, { nullptr, std::move(buff) });
+ tld.part[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status P_PARTKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_PARTKEY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_PARTKEY));
+ int32_t *p_partkey = reinterpret_cast<int32_t *>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.part_to_generate; i++)
+ {
+ p_partkey[i] = static_cast<int32_t>(tld.partkey_start + i + 1);
+ ARROW_DCHECK(1 <= p_partkey[i] && p_partkey[i] <= part_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_NAME(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_NAME].kind() == Datum::NONE)
+ {
+ std::uniform_int_distribution<int32_t> dist(0, static_cast<uint8_t>(kNumNameParts - 1));
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff, AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ size_t string_length = 0;
+ for(int ipart = 0; ipart < 5; ipart++)
+ {
+ uint8_t name_part_index = static_cast<uint8_t>(dist(tld.rng));
+ tld.string_indices[irow * 5 + ipart] = name_part_index;
+ string_length += std::strlen(NameParts[name_part_index]);
+ }
+ // Add 4 because there is a space between each word (i.e. four spaces)
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length + 4);
+ }
+ // Add an extra byte for the space after in the very last string.
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer, AllocateBuffer(offsets[tld.part_to_generate] + 1));
+ char *strings = reinterpret_cast<char *>(string_buffer->mutable_data());
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ char *row = strings + offsets[irow];
+ for(int ipart = 0; ipart < 5; ipart++)
+ {
+ uint8_t name_part_index = tld.string_indices[irow * 5 + ipart];
+ const char *part = NameParts[name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ *row++ = ' ';
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_NAME], tld.part_to_generate, { nullptr, std::move(offset_buff), std::move(string_buffer) });
+ Datum datum(ad);
+ tld.part[PART::P_NAME] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_MFGR(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_MFGR].kind() == Datum::NONE)
+ {
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char *manufacturer = "Manufacturer#";
+ const size_t manufacturer_length = std::strlen(manufacturer);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_MFGR));
+ char *p_mfgr = reinterpret_cast<char *>(tld.part[PART::P_MFGR].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ std::strncpy(p_mfgr + byte_width * irow, manufacturer, byte_width);
+ char mfgr_id = '0' + dist(tld.rng);
+ *(p_mfgr + byte_width * irow + manufacturer_length) = mfgr_id;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_BRAND(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_BRAND].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(P_MFGR(thread_index));
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char *brand = "Brand#";
+ const size_t brand_length = std::strlen(brand);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_BRAND));
+ const char *p_mfgr = reinterpret_cast<const char *>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->data());
+ char *p_brand = reinterpret_cast<char *>(
+ tld.part[PART::P_BRAND].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_BRAND]);
+ int32_t mfgr_byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ const size_t mfgr_id_offset = std::strlen("Manufacturer#");
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ char *row = p_brand + byte_width * irow;
+ char mfgr_id = *(p_mfgr + irow * mfgr_byte_width + mfgr_id_offset);
+ char brand_id = '0' + dist(tld.rng);
+ std::strncpy(row, brand, byte_width);
+ *(row + brand_length) = mfgr_id;
+ *(row + brand_length + 1) = brand_id;
+ irow += 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_TYPE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_TYPE].kind() == Datum::NONE)
+ {
+ using D = std::uniform_int_distribution<uint32_t>;
+ D dists[] =
+ {
+ D{ 0, static_cast<uint8_t>(kNumTypes_1 - 1) },
+ D{ 0, static_cast<uint8_t>(kNumTypes_2 - 1) },
+ D{ 0, static_cast<uint8_t>(kNumTypes_3 - 1) },
+ };
+
+ const char **types[] = { Types_1, Types_2, Types_3 };
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff, AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ size_t string_length = 0;
+ for(int ipart = 0; ipart < 3; ipart++)
+ {
+ uint8_t name_part_index = static_cast<uint8_t>(dists[ipart](tld.rng));
+ tld.string_indices[irow * 3 + ipart] = name_part_index;
+ string_length += std::strlen(types[ipart][name_part_index]);
+ }
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer, AllocateBuffer(offsets[tld.part_to_generate]));
+ char *strings = reinterpret_cast<char *>(string_buffer->mutable_data());
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ char *row = strings + offsets[irow];
+ for(int ipart = 0; ipart < 3; ipart++)
+ {
+ uint8_t name_part_index = tld.string_indices[irow * 3 + ipart];
+ const char *part = types[ipart][name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_TYPE], tld.part_to_generate, { nullptr, std::move(offset_buff), std::move(string_buffer) });
+ Datum datum(ad);
+ tld.part[PART::P_TYPE] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_SIZE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_SIZE].kind() == Datum::NONE)
+ {
+ std::uniform_int_distribution<int32_t> dist(1, 50);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_SIZE));
+ int32_t *p_size = reinterpret_cast<int32_t *>(
+ tld.part[PART::P_SIZE].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.part_to_generate; i++)
+ p_size[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status P_CONTAINER(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_CONTAINER].kind() == Datum::NONE)
+ {
+ std::uniform_int_distribution<int> dist1(0, static_cast<uint8_t>(kNumContainers_1 - 1));
+ std::uniform_int_distribution<int> dist2(0, static_cast<uint8_t>(kNumContainers_2 - 1));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_CONTAINER));
+ char *p_container = reinterpret_cast<char *>(
+ tld.part[PART::P_CONTAINER].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_CONTAINER]);
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ int container1_idx = dist1(tld.rng);
+ int container2_idx = dist2(tld.rng);
+ const char *container1 = Containers_1[container1_idx];
+ const char *container2 = Containers_2[container2_idx];
+ size_t container1_length = std::strlen(container1);
+ size_t container2_length = std::strlen(container2);
+
+ char *row = p_container + byte_width * irow;
+ std::strncpy(row, container1, byte_width);
+ std::memcpy(row + container1_length, container2, container2_length);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_RETAILPRICE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_RETAILPRICE].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_RETAILPRICE));
+ const int32_t *p_partkey = reinterpret_cast<const int32_t *>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+ Decimal128 *p_retailprice = reinterpret_cast<Decimal128 *>(
+ tld.part[PART::P_RETAILPRICE].array()->buffers[1]->mutable_data());
+ for(int64_t irow = 0; irow < tld.part_to_generate; irow++)
+ {
+ int32_t partkey = p_partkey[irow];
+ int64_t retail_price = (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ p_retailprice[irow] = { retail_price };
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_COMMENT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PART::P_COMMENT].kind() == Datum::NONE)
+ {
+ ARROW_ASSIGN_OR_RAISE(tld.part[PART::P_COMMENT], g_text.GenerateComments(tld.part_to_generate, 5, 22, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ int64_t PartsuppBatchesToGenerate(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t num_batches = (ps_to_generate + batch_size_ - 1) / batch_size_;
+ return num_batches;
+ }
+
+ Status InitPartsupp(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ tld.generated_partsupp.reset();
+ int64_t num_batches = PartsuppBatchesToGenerate(thread_index);
+ tld.partsupp.resize(num_batches);
+ for(std::vector<Datum> &batch : tld.partsupp)
+ {
+ batch.resize(PARTSUPP::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocatePartSuppBatch(size_t thread_index, size_t ibatch, int column)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartsuppTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff, AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kPartsuppTypes[column], batch_size_, { nullptr, std::move(buff) });
+ tld.partsupp[ibatch][column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status PS_PARTKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_partsupp[PARTSUPP::PS_PARTKEY])
+ {
+ tld.generated_partsupp[PARTSUPP::PS_PARTKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t *p_partkey = reinterpret_cast<const int32_t *>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for(int64_t irow = 0; irow < ps_to_generate; ibatch++)
+ {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_PARTKEY));
+ int32_t *ps_partkey = reinterpret_cast<int32_t *>(
+ tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY].array()->buffers[1]->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; ipartsupp < kPartSuppRowsPerPart && irun < next_run; ipartsupp++, irun++)
+ ps_partkey[batch_offset++] = p_partkey[ipart];
+
+ if(ipartsupp == kPartSuppRowsPerPart)
+ {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_partsupp[PARTSUPP::PS_SUPPKEY])
+ {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t *p_partkey = reinterpret_cast<const int32_t *>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for(int64_t irow = 0; irow < ps_to_generate; ibatch++)
+ {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPKEY));
+ int32_t *ps_suppkey = reinterpret_cast<int32_t *>(
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY].array()->buffers[1]->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; ipartsupp < kPartSuppRowsPerPart && irun < next_run; ipartsupp++, irun++)
+ {
+ int32_t supplier = static_cast<int32_t>(ipartsupp);
+ int32_t partkey = p_partkey[ipart];
+ ps_suppkey[batch_offset++] = (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ if(ipartsupp == kPartSuppRowsPerPart)
+ {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_AVAILQTY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_partsupp[PARTSUPP::PS_AVAILQTY])
+ {
+ tld.generated_partsupp[PARTSUPP::PS_AVAILQTY] = true;
+ std::uniform_int_distribution<int32_t> dist(1, 9999);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for(int64_t irow = 0; irow < ps_to_generate; ibatch++)
+ {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_AVAILQTY));
+ int32_t *ps_availqty = reinterpret_cast<int32_t *>(
+ tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY].array()->buffers[1]->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for(int64_t irun = 0; irun < next_run; irun++)
+ ps_availqty[irun] = dist(tld.rng);
+
+ tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPLYCOST(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST])
+ {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST] = true;
+ std::uniform_int_distribution<int64_t> dist(100, 100000);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for(int64_t irow = 0; irow < ps_to_generate; ibatch++)
+ {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPLYCOST));
+ Decimal128 *ps_supplycost = reinterpret_cast<Decimal128 *>(
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST].array()->buffers[1]->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for(int64_t irun = 0; irun < next_run; irun++)
+ ps_supplycost[irun] = { dist(tld.rng) };
+
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_COMMENT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.part[PARTSUPP::PS_COMMENT].kind() == Datum::NONE)
+ {
+ int64_t irow = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for(size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++)
+ {
+ int64_t num_rows = std::min(batch_size_, ps_to_generate - irow);
+ ARROW_ASSIGN_OR_RAISE(
+ tld.partsupp[ibatch][PARTSUPP::PS_COMMENT], g_text.GenerateComments(num_rows, 49, 198, tld.rng));
+ irow += num_rows;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData
+ {
+ std::vector<Datum> part;
+ std::vector<int8_t> string_indices;
+ int64_t part_to_generate{0};
+ int64_t partkey_start{0};
+
+ std::vector<std::vector<Datum>> partsupp;
+ std::bitset<PARTSUPP::kNumCols> generated_partsupp;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex part_output_queue_mutex_;
+ std::mutex partsupp_output_queue_mutex_;
+ std::queue<ExecBatch> part_output_queue_;
+ std::queue<ExecBatch> partsupp_output_queue_;
+ int64_t batch_size_{0};
+ float scale_factor_{0};
+ int64_t part_rows_to_generate_{0};
+ int64_t part_rows_generated_{0};
+ std::vector<int> part_cols_;
+ std::vector<int> partsupp_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<int64_t> part_batches_generated_ = { 0 };
+ std::atomic<int64_t> partsupp_batches_generated_ = { 0 };
+ static constexpr int64_t kPartSuppRowsPerPart = 4;
+ };
+
+ class OrdersAndLineItemGenerator
+ {
+ public:
+ Status Init(
+ size_t num_threads,
+ int64_t batch_size,
+ float scale_factor)
+ {
+ if(!inited_)
+ {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for(ThreadLocalData &tld : thread_local_data_)
+ {
+ tld.items_per_order.resize(batch_size_);
+ tld.rng.seed(seq);
+ }
+ orders_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000 * 10);
+ }
+ return Status::OK();
+ }
+
+ int64_t orders_batches_generated() const
+ {
+ return orders_batches_generated_.load();
+ }
+
+ int64_t lineitem_batches_generated() const
+ {
+ return lineitem_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetOrdersOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, kOrdersTypes, kOrdersNameMap, orders_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetLineItemOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, kLineitemTypes, kLineitemNameMap, lineitem_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextOrdersBatch()
+ {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if(!orders_output_queue_.empty())
+ {
+ ExecBatch batch = std::move(orders_output_queue_.front());
+ orders_output_queue_.pop();
+ return std::move(batch);
+ }
+ else if(orders_rows_generated_ == orders_rows_to_generate_)
+ {
+ return util::nullopt;
+ }
+ else
+ {
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate = std::min(
+ batch_size_,
+ orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.first_batch_offset = 0;
+ tld.generated_lineitem.reset();
+
+ for(int col : orders_cols_)
+ RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for(int col : lineitem_cols_)
+ RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for(size_t i = 0; i < orders_cols_.size(); i++)
+ {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ if(!lineitem_cols_.empty())
+ {
+ std::vector<ExecBatch> lineitem_results;
+ for(size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++)
+ {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for(size_t icol = 0; icol < lineitem_cols_.size(); icol++)
+ {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(lineitem_output_queue_mutex_);
+ for(ExecBatch &eb : lineitem_results)
+ {
+ lineitem_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(orders_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextLineItemBatch()
+ {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ ExecBatch queued;
+ bool from_queue = false;
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ if(!lineitem_output_queue_.empty())
+ {
+ queued = std::move(lineitem_output_queue_.front());
+ lineitem_output_queue_.pop();
+ from_queue = true;
+ }
+ }
+ tld.first_batch_offset = 0;
+ if(from_queue)
+ {
+ ARROW_DCHECK(queued.length <= batch_size_);
+ tld.first_batch_offset = queued.length;
+ if(queued.length == batch_size_)
+ return std::move(queued);
+ }
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if(orders_rows_generated_ == orders_rows_to_generate_)
+ {
+ if(from_queue)
+ return std::move(queued);
+ return util::nullopt;
+ }
+
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate = std::min(
+ batch_size_,
+ orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1ll);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.generated_lineitem.reset();
+ if(from_queue)
+ {
+ lineitem_batches_generated_.fetch_sub(1);
+ for(size_t i = 0; i < lineitem_cols_.size(); i++)
+ if(tld.lineitem[0][lineitem_cols_[i]].kind() == Datum::NONE)
+ tld.lineitem[0][lineitem_cols_[i]] = std::move(queued[i]);
+ }
+
+ for(int col : orders_cols_)
+ RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for(int col : lineitem_cols_)
+ RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ if(!orders_cols_.empty())
+ {
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for(size_t i = 0; i < orders_cols_.size(); i++)
+ {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch orders_batch, ExecBatch::Make(std::move(orders_result)));
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ orders_output_queue_.emplace(std::move(orders_batch));
+ }
+ }
+ std::vector<ExecBatch> lineitem_results;
+ for(size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++)
+ {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for(size_t icol = 0; icol < lineitem_cols_.size(); icol++)
+ {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ lineitem_batches_generated_.fetch_add(static_cast<int64_t>(lineitem_results.size()));
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ for(size_t i = 1; i < lineitem_results.size(); i++)
+ lineitem_output_queue_.emplace(std::move(lineitem_results[i]));
+ }
+ return std::move(lineitem_results[0]);
+ }
+
+ private:
+#define FOR_EACH_ORDERS_COLUMN(F) \
+ F(O_ORDERKEY) \
+ F(O_CUSTKEY) \
+ F(O_ORDERSTATUS) \
+ F(O_TOTALPRICE) \
+ F(O_ORDERDATE) \
+ F(O_ORDERPRIORITY) \
+ F(O_CLERK) \
+ F(O_SHIPPRIORITY) \
+ F(O_COMMENT)
+
+#define FOR_EACH_LINEITEM_COLUMN(F) \
+ F(L_ORDERKEY) \
+ F(L_PARTKEY) \
+ F(L_SUPPKEY) \
+ F(L_LINENUMBER) \
+ F(L_QUANTITY) \
+ F(L_EXTENDEDPRICE) \
+ F(L_DISCOUNT) \
+ F(L_TAX) \
+ F(L_RETURNFLAG) \
+ F(L_LINESTATUS) \
+ F(L_SHIPDATE) \
+ F(L_COMMITDATE) \
+ F(L_RECEIPTDATE) \
+ F(L_SHIPINSTRUCT) \
+ F(L_SHIPMODE) \
+ F(L_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct ORDERS
+ {
+ enum
+ {
+ FOR_EACH_ORDERS_COLUMN(MAKE_ENUM)
+ kNumCols,
+ };
+ };
+ struct LINEITEM
+ {
+ enum
+ {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_ENUM)
+ kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) \
+ { #col, ORDERS::col },
+ const std::unordered_map<std::string, int> kOrdersNameMap =
+ {
+ FOR_EACH_ORDERS_COLUMN(MAKE_STRING_MAP)
+ };
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) \
+ { #col, LINEITEM::col },
+ const std::unordered_map<std::string, int> kLineitemNameMap =
+ {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_STRING_MAP)
+ };
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ const std::vector<GenerateColumnFn> kOrdersGenerators =
+ {
+ FOR_EACH_ORDERS_COLUMN(MAKE_FN_ARRAY)
+ };
+ const std::vector<GenerateColumnFn> kLineitemGenerators =
+ {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_FN_ARRAY)
+ };
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kOrdersTypes =
+ {
+ int32(),
+ int32(),
+ fixed_size_binary(1),
+ decimal(12, 2),
+ date32(),
+ fixed_size_binary(15),
+ fixed_size_binary(15),
+ int32(),
+ utf8()
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kLineitemTypes =
+ {
+ int32(),
+ int32(),
+ int32(),
+ int32(),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ fixed_size_binary(1),
+ fixed_size_binary(1),
+ date32(),
+ date32(),
+ date32(),
+ fixed_size_binary(25),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status AllocateOrdersBatch(size_t thread_index, int column)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.orders[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff, AllocateBuffer(tld.orders_to_generate * byte_width));
+ ArrayData ad(kOrdersTypes[column], tld.orders_to_generate, { nullptr, std::move(buff) });
+ tld.orders[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status O_ORDERKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_ORDERKEY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERKEY));
+ int32_t *o_orderkey = reinterpret_cast<int32_t *>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.orders_to_generate; i++)
+ {
+ int32_t orderkey_index = static_cast<int32_t>(tld.orderkey_start + i);
+ int32_t index_of_run = orderkey_index / 8;
+ int32_t index_in_run = orderkey_index % 8;
+ o_orderkey[i] = (index_of_run * 32 + index_in_run + 1);
+ ARROW_DCHECK(1 <= o_orderkey[i] && o_orderkey[i] <= 4 * orders_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CUSTKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_CUSTKEY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CUSTKEY));
+
+ // Spec says it must be a random number between 1 and SF*150000 that is not
+ // divisible by 3. Rather than repeatedly generating numbers until we get to
+ // a non-divisible-by-3 number, we just generate a number between
+ // 0 and SF * 50000 - 1, multiply by 3, and then add either 1 or 2.
+ int32_t sf_50k = static_cast<int32_t>(scale_factor_ * 50000);
+ std::uniform_int_distribution<int32_t> base_dist(0, sf_50k - 1);
+ std::uniform_int_distribution<int32_t> offset_dist(1, 2);
+ int32_t *o_custkey = reinterpret_cast<int32_t *>(
+ tld.orders[ORDERS::O_CUSTKEY].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.orders_to_generate; i++)
+ o_custkey[i] = 3 * base_dist(tld.rng) + offset_dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERSTATUS(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_ORDERSTATUS].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(L_LINESTATUS(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERSTATUS));
+
+ char *o_orderstatus = reinterpret_cast<char *>(
+ tld.orders[ORDERS::O_ORDERSTATUS].array()->buffers[1]->mutable_data());
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ bool all_f = true;
+ bool all_o = true;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ const char *l_linestatus = reinterpret_cast<const char *>(
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->buffers[1]->data());
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++, batch_offset++)
+ {
+ all_f &= l_linestatus[batch_offset] == 'F';
+ all_o &= l_linestatus[batch_offset] == 'O';
+ }
+ if(iline == tld.items_per_order[iorder])
+ {
+ iline = 0;
+ ARROW_DCHECK(!(all_f && all_o));
+ if(all_f)
+ o_orderstatus[iorder] = 'F';
+ else if(all_o)
+ o_orderstatus[iorder] = 'O';
+ else
+ o_orderstatus[iorder] = 'P';
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_TOTALPRICE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_TOTALPRICE].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(L_EXTENDEDPRICE(thread_index));
+ RETURN_NOT_OK(L_TAX(thread_index));
+ RETURN_NOT_OK(L_DISCOUNT(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_TOTALPRICE));
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ int64_t sum = 0;
+ Decimal128 *o_totalprice = reinterpret_cast<Decimal128 *>(
+ tld.orders[ORDERS::O_TOTALPRICE].array()->buffers[1]->mutable_data());
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ const Decimal128 *l_extendedprice = reinterpret_cast<const Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->buffers[1]->data());
+ const Decimal128 *l_tax = reinterpret_cast<const Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->data());
+ const Decimal128 *l_discount = reinterpret_cast<const Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->buffers[1]->data());
+
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++, batch_offset++)
+ {
+ int64_t eprice = static_cast<int64_t>(l_extendedprice[batch_offset]);
+ int64_t tax = static_cast<int64_t>(l_tax[batch_offset]);
+ int64_t discount = static_cast<int64_t>(l_discount[batch_offset]);
+ sum += (eprice * (100 + tax) * (100 - discount));
+ }
+ if(iline == tld.items_per_order[iorder])
+ {
+ sum /= 100 * 100;
+ o_totalprice[iorder] = { sum };
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERDATE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_ORDERDATE].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERDATE));
+
+ std::uniform_int_distribution<uint32_t> dist(kStartDate, kEndDate - 151);
+ uint32_t *o_orderdate = reinterpret_cast<uint32_t *>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.orders_to_generate; i++)
+ o_orderdate[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERPRIORITY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_ORDERPRIORITY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERPRIORITY));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_ORDERPRIORITY]);
+ std::uniform_int_distribution<int32_t> dist(0, kNumPriorities - 1);
+ char *o_orderpriority = reinterpret_cast<char *>(
+ tld.orders[ORDERS::O_ORDERPRIORITY].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.orders_to_generate; i++)
+ {
+ const char *str = Priorities[dist(tld.rng)];
+ std::strncpy(o_orderpriority + i * byte_width, str, byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CLERK(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_CLERK].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CLERK));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_CLERK]);
+ int64_t max_clerk_id = static_cast<int64_t>(scale_factor_ * 1000);
+ std::uniform_int_distribution<int64_t> dist(1, max_clerk_id);
+ char *o_clerk = reinterpret_cast<char *>(
+ tld.orders[ORDERS::O_CLERK].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < tld.orders_to_generate; i++)
+ {
+ const char *clerk = "Clerk#";
+ const size_t clerk_length = std::strlen(clerk);
+ int64_t clerk_number = dist(tld.rng);
+ char *output = o_clerk + i * byte_width;
+ std::strncpy(output, clerk, byte_width);
+ AppendNumberPaddedToNineDigits(output + clerk_length, clerk_number);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_SHIPPRIORITY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_SHIPPRIORITY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_SHIPPRIORITY));
+ int32_t *o_shippriority = reinterpret_cast<int32_t *>(
+ tld.orders[ORDERS::O_SHIPPRIORITY].array()->buffers[1]->mutable_data());
+ std::memset(o_shippriority, 0, tld.orders_to_generate * sizeof(int32_t));
+ }
+ return Status::OK();
+ }
+
+ Status O_COMMENT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.orders[ORDERS::O_COMMENT].kind() == Datum::NONE)
+ {
+ ARROW_ASSIGN_OR_RAISE(tld.orders[ORDERS::O_COMMENT], g_text.GenerateComments(tld.orders_to_generate, 19, 78, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ Status GenerateRowCounts(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ std::uniform_int_distribution<int> length_dist(1, 7);
+ tld.lineitem_to_generate = 0;
+ tld.items_per_order.clear();
+ for(int64_t i = 0; i < tld.orders_to_generate; i++)
+ {
+ int length = length_dist(tld.rng);
+ tld.items_per_order.push_back(length);
+ tld.lineitem_to_generate += length;
+ }
+ int64_t num_batches = (tld.first_batch_offset + tld.lineitem_to_generate + batch_size_ - 1) / batch_size_;
+ tld.lineitem.resize(num_batches);
+ for(std::vector<Datum> &batch : tld.lineitem)
+ {
+ batch.resize(LINEITEM::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocateLineItemBufferIfNeeded(size_t thread_index, size_t ibatch, int column, size_t &out_batch_offset)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.lineitem[ibatch][column].kind() == Datum::NONE)
+ {
+ int32_t byte_width = arrow::internal::GetByteWidth(*kLineitemTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff, AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kLineitemTypes[column], batch_size_, { nullptr, std::move(buff) });
+ tld.lineitem[ibatch][column] = std::move(ad);
+ out_batch_offset = 0;
+ }
+ if(ibatch == 0)
+ out_batch_offset = tld.first_batch_offset;
+
+ return Status::OK();
+ }
+
+ Status L_ORDERKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_ORDERKEY])
+ {
+ tld.generated_lineitem[LINEITEM::L_ORDERKEY] = true;
+ RETURN_NOT_OK(O_ORDERKEY(thread_index));
+ const int32_t *o_orderkey = reinterpret_cast<const int32_t *>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_ORDERKEY, batch_offset));
+ int32_t *l_linenumber = reinterpret_cast<int32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_ORDERKEY].array()->buffers[1]->mutable_data());
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_linenumber[batch_offset++] = o_orderkey[iorder];
+ if(iline == tld.items_per_order[iorder])
+ {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_ORDERKEY].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_PARTKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_PARTKEY])
+ {
+ tld.generated_lineitem[LINEITEM::L_PARTKEY] = true;
+
+ size_t ibatch = 0;
+ int32_t max_partkey = static_cast<int32_t>(scale_factor_ * 200000);
+ std::uniform_int_distribution<int32_t> dist(1, max_partkey);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_PARTKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t *l_partkey = reinterpret_cast<int32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_partkey[batch_offset] = dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SUPPKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_SUPPKEY])
+ {
+ tld.generated_lineitem[LINEITEM::L_SUPPKEY] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 3);
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset = 0;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_SUPPKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t *l_suppkey = reinterpret_cast<int32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_SUPPKEY].array()->buffers[1]->mutable_data());
+ const int32_t *l_partkey = reinterpret_cast<const int32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ for(int64_t i = 0; i < next_run; i++)
+ {
+ int32_t supplier = dist(tld.rng);
+ int32_t partkey = l_partkey[batch_offset];
+ // Fun fact: the parentheses for this expression are unbalanced in the TPC-H spec.
+ l_suppkey[batch_offset++] = (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SUPPKEY].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINENUMBER(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_LINENUMBER])
+ {
+ tld.generated_lineitem[LINEITEM::L_LINENUMBER] = true;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_LINENUMBER, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t *l_linenumber = reinterpret_cast<int32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_LINENUMBER].array()->buffers[1]->mutable_data());
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ {
+ l_linenumber[batch_offset++] = (iline + 1);
+ ARROW_DCHECK(1 <= (iline + 1) && (iline + 1) <= 7);
+ }
+ if(iline == tld.items_per_order[iorder])
+ {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINENUMBER].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_QUANTITY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_QUANTITY])
+ {
+ tld.generated_lineitem[LINEITEM::L_QUANTITY] = true;
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int64_t> dist(1, 50);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_QUANTITY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128 *l_quantity = reinterpret_cast<Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->buffers[1]->mutable_data());
+ for(int64_t i = 0; i < next_run; i++)
+ {
+ // Multiply by 100 because the type is decimal(12, 2), so the decimal goes after two digits
+ int64_t quantity = dist(tld.rng) * 100;
+ l_quantity[batch_offset++] = { quantity };
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_EXTENDEDPRICE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE])
+ {
+ tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+ RETURN_NOT_OK(L_QUANTITY(thread_index));
+ size_t ibatch = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_EXTENDEDPRICE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ const int32_t *l_partkey = reinterpret_cast<const int32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ const Decimal128 *l_quantity = reinterpret_cast<const Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->buffers[1]->data());
+ Decimal128 *l_extendedprice = reinterpret_cast<Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ {
+ int64_t partkey = static_cast<int64_t>(l_partkey[batch_offset]);
+ // Divide by 100 to recover the integer representation (not Decimal).
+ int64_t quantity = static_cast<int64_t>(l_quantity[batch_offset]) / 100;
+
+ // Spec says to divide by 100, but that happens automatically due to this being stored
+ // to two decimal points.
+ int64_t retail_price = (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ int64_t extended_price = retail_price * quantity;
+ l_extendedprice[batch_offset] = { extended_price };
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_DISCOUNT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_DISCOUNT])
+ {
+ tld.generated_lineitem[LINEITEM::L_DISCOUNT] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 10);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_DISCOUNT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128 *l_discount = reinterpret_cast<Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_discount[batch_offset] = { dist(tld.rng) };
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_TAX(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_TAX])
+ {
+ tld.generated_lineitem[LINEITEM::L_TAX] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 8);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_TAX, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128 *l_tax = reinterpret_cast<Decimal128 *>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_tax[batch_offset] = { dist(tld.rng) };
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RETURNFLAG(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_RETURNFLAG])
+ {
+ tld.generated_lineitem[LINEITEM::L_RETURNFLAG] = true;
+ RETURN_NOT_OK(L_RECEIPTDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<uint32_t> dist;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_RETURNFLAG, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char *l_returnflag = reinterpret_cast<char *>(
+ tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG].array()->buffers[1]->mutable_data());
+ const uint32_t *l_receiptdate = reinterpret_cast<const uint32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ {
+ if(l_receiptdate[batch_offset] <= kCurrentDate)
+ {
+ uint32_t r = dist(tld.rng);
+ l_returnflag[batch_offset] = (r % 2 == 1) ? 'R' : 'A';
+ }
+ else
+ {
+ l_returnflag[batch_offset] = 'N';
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINESTATUS(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_LINESTATUS])
+ {
+ tld.generated_lineitem[LINEITEM::L_LINESTATUS] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_LINESTATUS, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char *l_linestatus = reinterpret_cast<char *>(
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->buffers[1]->mutable_data());
+ const uint32_t *l_shipdate = reinterpret_cast<const uint32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ {
+ if(l_shipdate[batch_offset] > kCurrentDate)
+ l_linestatus[batch_offset] = 'O';
+ else
+ l_linestatus[batch_offset] = 'F';
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPDATE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_SHIPDATE])
+ {
+ tld.generated_lineitem[LINEITEM::L_SHIPDATE] = true;
+ RETURN_NOT_OK(O_ORDERDATE(thread_index));
+ const int32_t *o_orderdate = reinterpret_cast<const int32_t *>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(1, 121);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_SHIPDATE, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t *l_shipdate = reinterpret_cast<uint32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->buffers[1]->mutable_data());
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_shipdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if(iline == tld.items_per_order[iorder])
+ {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMITDATE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_COMMITDATE])
+ {
+ tld.generated_lineitem[LINEITEM::L_COMMITDATE] = true;
+ const int32_t *o_orderdate = reinterpret_cast<const int32_t *>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(30, 90);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_COMMITDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t *l_commitdate = reinterpret_cast<uint32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_COMMITDATE].array()->buffers[1]->mutable_data());
+ for(int64_t irun = 0; irun < next_run;)
+ {
+ for(; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_commitdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if(iline == tld.items_per_order[iorder])
+ {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMITDATE].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RECEIPTDATE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_RECEIPTDATE])
+ {
+ tld.generated_lineitem[LINEITEM::L_RECEIPTDATE] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(1, 30);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_RECEIPTDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t *l_receiptdate = reinterpret_cast<uint32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE].array()->buffers[1]->mutable_data());
+ const uint32_t *l_shipdate = reinterpret_cast<const uint32_t *>(
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_receiptdate[batch_offset] = l_shipdate[batch_offset] + dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPINSTRUCT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT])
+ {
+ tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT] = true;
+ int32_t byte_width = arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPINSTRUCT]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumInstructions - 1);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_SHIPINSTRUCT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char *l_shipinstruct = reinterpret_cast<char *>(
+ tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ {
+ const char *str = Instructions[dist(tld.rng)];
+ // Note that we don't have to memset the buffer to 0 because strncpy pads each string
+ // with 0's anyway
+ std::strncpy(l_shipinstruct + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPMODE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_SHIPMODE])
+ {
+ tld.generated_lineitem[LINEITEM::L_SHIPMODE] = true;
+ int32_t byte_width = arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPMODE]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumModes - 1);
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch, LINEITEM::L_SHIPMODE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char *l_shipmode = reinterpret_cast<char *>(
+ tld.lineitem[ibatch][LINEITEM::L_SHIPMODE].array()->buffers[1]->mutable_data());
+
+ for(int64_t i = 0; i < next_run; i++, batch_offset++)
+ {
+ const char *str = Modes[dist(tld.rng)];
+ std::strncpy(l_shipmode + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPMODE].array()->length = static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMENT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(!tld.generated_lineitem[LINEITEM::L_COMMENT])
+ {
+ tld.generated_lineitem[LINEITEM::L_COMMENT] = true;
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ for(int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++)
+ {
+ // Comments are kind of sneaky: we always generate the full batch and then just bump the length
+ if(tld.lineitem[ibatch][LINEITEM::L_COMMENT].kind() == Datum::NONE)
+ {
+ ARROW_ASSIGN_OR_RAISE(tld.lineitem[ibatch][LINEITEM::L_COMMENT], g_text.GenerateComments(batch_size_, 10, 43, tld.rng));
+ batch_offset = 0;
+ }
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ batch_offset += next_run;
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMENT].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData
+ {
+ std::vector<Datum> orders;
+ int64_t orders_to_generate;
+ int64_t orderkey_start;
+
+ std::vector<std::vector<Datum>> lineitem;
+ std::vector<int> items_per_order;
+ int64_t lineitem_to_generate;
+ int64_t first_batch_offset;
+ std::bitset<LINEITEM::kNumCols> generated_lineitem;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex orders_output_queue_mutex_;
+ std::mutex lineitem_output_queue_mutex_;
+ std::queue<ExecBatch> orders_output_queue_;
+ std::queue<ExecBatch> lineitem_output_queue_;
+ int64_t batch_size_;
+ float scale_factor_;
+ int64_t orders_rows_to_generate_;
+ int64_t orders_rows_generated_;
+ std::vector<int> orders_cols_;
+ std::vector<int> lineitem_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<size_t> orders_batches_generated_ = { 0 };
+ std::atomic<size_t> lineitem_batches_generated_ = { 0 };
+ };
+
+ class SupplierGenerator : public TpchTableGenerator
+ {
+ public:
+ Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) override
+ {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 10000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_, SetOutputColumns(
+ columns,
+ kTypes,
+ kNameMap,
+ gen_list_));
+
+ random::pcg32_fast rng;
+ std::uniform_int_distribution<int64_t> dist(0, rows_to_generate_ - 1);
+ size_t num_special_rows = static_cast<size_t>(5 * scale_factor_);
+ std::unordered_set<int64_t> good_rows_set;
+ while(good_rows_set.size() < num_special_rows)
+ {
+ int64_t row = dist(rng);
+ good_rows_set.insert(row);
+ }
+ std::unordered_set<int64_t> bad_rows_set;
+ while(bad_rows_set.size() < num_special_rows)
+ {
+ int64_t bad_row;
+ do
+ {
+ bad_row = dist(rng);
+ } while(good_rows_set.find(bad_row) != good_rows_set.end());
+ bad_rows_set.insert(bad_row);
+ }
+ good_rows_.clear();
+ bad_rows_.clear();
+ good_rows_.insert(good_rows_.end(), good_rows_set.begin(), good_rows_set.end());
+ bad_rows_.insert(bad_rows_.end(), bad_rows_set.begin(), bad_rows_set.end());
+ std::sort(good_rows_.begin(), good_rows_.end());
+ std::sort(bad_rows_.begin(), bad_rows_.end());
+ return Status::OK();
+ }
+
+ Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override
+ {
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for(ThreadLocalData &tld : thread_local_data_)
+ tld.rng.seed(seq);
+
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for(size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override
+ {
+ return schema_;
+ }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(S_SUPPKEY) \
+ F(S_NAME) \
+ F(S_ADDRESS) \
+ F(S_NATIONKEY) \
+ F(S_PHONE) \
+ F(S_ACCTBAL) \
+ F(S_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct SUPPLIER
+ {
+ enum
+ {
+ FOR_EACH_COLUMN(MAKE_ENUM)
+ kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) \
+ { #col, SUPPLIER::col },
+ const std::unordered_map<std::string, int> kNameMap =
+ {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)
+ };
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kGenerators =
+ {
+ FOR_EACH_COLUMN(MAKE_FN_ARRAY)
+ };
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> kTypes =
+ {
+ int32(),
+ fixed_size_binary(25),
+ utf8(),
+ int32(),
+ fixed_size_binary(15),
+ decimal(12, 2),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index)
+ {
+ if(done_.load())
+ return Status::OK();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ tld.suppkey_start = rows_generated_.fetch_add(batch_size_);
+ if(tld.suppkey_start >= rows_to_generate_)
+ return Status::OK();
+
+ tld.to_generate = std::min(batch_size_,
+ rows_to_generate_ - tld.suppkey_start);
+
+ tld.batch.resize(SUPPLIER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for(int col : gen_list_)
+ RETURN_NOT_OK(kGenerators[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for(size_t i = 0; i < gen_list_.size(); i++)
+ {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_outputted_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_outputted_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if(is_last_batch)
+ {
+ done_.store(true);
+ finished_callback_(batches_outputted_.load());
+ return Status::OK();
+ }
+ return schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff, AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(kTypes[column], tld.to_generate, { nullptr, std::move(buff) });
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status S_SUPPKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[SUPPLIER::S_SUPPKEY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_SUPPKEY));
+ int32_t *s_suppkey = reinterpret_cast<int32_t *>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->mutable_data());
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ {
+ s_suppkey[irow] = static_cast<int32_t>(tld.suppkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_NAME(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[SUPPLIER::S_NAME].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(S_SUPPKEY(thread_index));
+ const int32_t *s_suppkey = reinterpret_cast<const int32_t *>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->data());
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NAME));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[SUPPLIER::S_NAME]);
+ char *s_name = reinterpret_cast<char *>(
+ tld.batch[SUPPLIER::S_NAME].array()->buffers[1]->mutable_data());
+ // Look man, I'm just following the spec ok? Section 4.2.3 as of March 1 2022
+ const char *supplier = "Supplie#r";
+ const size_t supplier_length = std::strlen(supplier);
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ {
+ char *out = s_name + byte_width * irow;
+ std::strncpy(out, supplier, byte_width);
+ AppendNumberPaddedToNineDigits(out + supplier_length, s_suppkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ADDRESS(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[SUPPLIER::S_ADDRESS].kind() == Datum::NONE)
+ {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.batch[SUPPLIER::S_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status S_NATIONKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[SUPPLIER::S_NATIONKEY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t *s_nationkey = reinterpret_cast<int32_t *>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->mutable_data());
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status S_PHONE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[SUPPLIER::S_PHONE].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(S_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[SUPPLIER::S_PHONE]);
+ const int32_t *s_nationkey = reinterpret_cast<const int32_t *>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->data());
+ char *s_phone = reinterpret_cast<char *>(
+ tld.batch[SUPPLIER::S_PHONE].array()->buffers[1]->mutable_data());
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ {
+ GeneratePhoneNumber(
+ s_phone + irow * byte_width,
+ tld.rng,
+ s_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ACCTBAL(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[SUPPLIER::S_ACCTBAL].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_ACCTBAL));
+ Decimal128 *s_acctbal = reinterpret_cast<Decimal128 *>(
+ tld.batch[SUPPLIER::S_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_acctbal[irow] = { dist(tld.rng) };
+ }
+ return Status::OK();
+ }
+
+ Status S_COMMENT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[SUPPLIER::S_COMMENT].kind() == Datum::NONE)
+ {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[SUPPLIER::S_COMMENT], g_text.GenerateComments(tld.to_generate, 25, 100, tld.rng));
+ ModifyComments(thread_index, "Recommends", good_rows_);
+ ModifyComments(thread_index, "Complaints", bad_rows_);
+ }
+ return Status::OK();
+ }
+
+ void ModifyComments(
+ size_t thread_index,
+ const char *review,
+ const std::vector<int64_t> &indices)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ const int32_t *offsets = reinterpret_cast<const int32_t *>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[1]->data());
+ char *str = reinterpret_cast<char *>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[2]->mutable_data());
+ const char *customer = "Customer";
+ const int32_t customer_length = static_cast<int32_t>(std::strlen(customer));
+ const int32_t review_length = static_cast<int32_t>(std::strlen(review));
+
+ auto it = std::lower_bound(indices.begin(), indices.end(), tld.suppkey_start);
+ for(; it != indices.end() && *it < tld.suppkey_start + tld.to_generate; it++)
+ {
+ int64_t idx_in_batch = *it - tld.suppkey_start;
+ char *out = str + offsets[idx_in_batch];
+ int32_t str_length = offsets[idx_in_batch + 1] - offsets[idx_in_batch];
+ std::uniform_int_distribution<int32_t> gap_dist(0, str_length - customer_length - review_length);
+ int32_t gap = gap_dist(tld.rng);
+ int32_t total_length = customer_length + gap + review_length;
+ std::uniform_int_distribution<int32_t> start_dist(0, str_length - total_length);
+ int32_t start = start_dist(tld.rng);
+ std::memcpy(out + start, customer, customer_length);
+ std::memcpy(out + start + customer_length + gap, review, review_length);
+ }
+ }
+
+ struct ThreadLocalData
+ {
+ random::pcg32_fast rng;
+ int64_t suppkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+ std::vector<int64_t> good_rows_;
+ std::vector<int64_t> bad_rows_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_;
+ std::atomic<int64_t> rows_generated_;
+ float scale_factor_;
+ int64_t batch_size_;
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+ };
+
+ class PartGenerator : public TpchTableGenerator
+ {
+ public:
+ PartGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen))
+ {
+ }
+
+ Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) override
+ {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ gen_->SetPartOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override
+ {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for(size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override
+ {
+ return schema_;
+ }
+
+ private:
+ Status ProduceCallback(size_t)
+ {
+ if(done_.load())
+ return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextPartBatch());
+ if(!maybe_batch.has_value())
+ {
+ int64_t batches_generated = gen_->part_batches_generated();
+ if(batches_generated == batches_outputted_.load())
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+ };
+
+ class PartSuppGenerator : public TpchTableGenerator
+ {
+ public:
+ PartSuppGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen))
+ {
+ }
+
+ Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) override
+ {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ gen_->SetPartSuppOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override
+ {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for(size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override
+ {
+ return schema_;
+ }
+
+ private:
+ Status ProduceCallback(size_t)
+ {
+ if(done_.load())
+ return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextPartSuppBatch());
+ if(!maybe_batch.has_value())
+ {
+ int64_t batches_generated = gen_->partsupp_batches_generated();
+ if(batches_generated == batches_outputted_.load())
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+ };
+
+ class CustomerGenerator : public TpchTableGenerator
+ {
+ public:
+ Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) override
+ {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_, SetOutputColumns(
+ columns,
+ kTypes,
+ kNameMap,
+ gen_list_));
+ return Status::OK();
+ }
+
+ Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override
+ {
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for(ThreadLocalData &tld : thread_local_data_)
+ tld.rng.seed(seq);
+
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for(size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override
+ {
+ return schema_;
+ }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(C_CUSTKEY) \
+ F(C_NAME) \
+ F(C_ADDRESS) \
+ F(C_NATIONKEY) \
+ F(C_PHONE) \
+ F(C_ACCTBAL) \
+ F(C_MKTSEGMENT) \
+ F(C_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct CUSTOMER
+ {
+ enum
+ {
+ FOR_EACH_COLUMN(MAKE_ENUM)
+ kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) \
+ { #col, CUSTOMER::col },
+ const std::unordered_map<std::string, int> kNameMap =
+ {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)
+ };
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kGenerators =
+ {
+ FOR_EACH_COLUMN(MAKE_FN_ARRAY)
+ };
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> kTypes =
+ {
+ int32(),
+ utf8(),
+ utf8(),
+ int32(),
+ fixed_size_binary(15),
+ decimal(12, 2),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index)
+ {
+ if(done_.load())
+ return Status::OK();
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ tld.custkey_start = rows_generated_.fetch_add(batch_size_);
+ if(tld.custkey_start >= rows_to_generate_)
+ return Status::OK();
+
+ tld.to_generate = std::min(batch_size_,
+ rows_to_generate_ - tld.custkey_start);
+
+ tld.batch.resize(CUSTOMER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for(int col : gen_list_)
+ RETURN_NOT_OK(kGenerators[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for(size_t i = 0; i < gen_list_.size(); i++)
+ {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_generated_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_generated_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if(is_last_batch)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ return schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff, AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(kTypes[column], tld.to_generate, { nullptr, std::move(buff) });
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status C_CUSTKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[CUSTOMER::C_CUSTKEY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_CUSTKEY));
+ int32_t *c_custkey = reinterpret_cast<int32_t *>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->mutable_data());
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ {
+ c_custkey[irow] = static_cast<int32_t>(tld.custkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_NAME(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[CUSTOMER::C_NAME].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(C_CUSTKEY(thread_index));
+ const int32_t *c_custkey = reinterpret_cast<const int32_t *>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->data());
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff, AllocateBuffer((tld.to_generate + 1) * sizeof(int32_t)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buff->mutable_data());
+ const char *customer = "Customer#";
+ const size_t customer_length = std::strlen(customer);
+ offsets[0] = 0;
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ {
+ int num_digits = GetNumDigits(c_custkey[irow]);
+ int num_chars = std::max(num_digits, 9);
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + num_chars + customer_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff, AllocateBuffer(offsets[tld.to_generate]));
+ char *str = reinterpret_cast<char *>(str_buff->mutable_data());
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ {
+ char *out = str + offsets[irow];
+ std::memcpy(out, customer, customer_length);
+ AppendNumberPaddedToNineDigits(out + customer_length, c_custkey[irow]);
+ }
+ ArrayData ad(utf8(), tld.to_generate, { nullptr, std::move(offset_buff), std::move(str_buff) });
+ tld.batch[CUSTOMER::C_NAME] = std::move(ad);
+ }
+ return Status::OK();
+ }
+
+ Status C_ADDRESS(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[CUSTOMER::C_ADDRESS].kind() == Datum::NONE)
+ {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.batch[CUSTOMER::C_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status C_NATIONKEY(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[CUSTOMER::C_NATIONKEY].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t *c_nationkey = reinterpret_cast<int32_t *>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->mutable_data());
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status C_PHONE(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[CUSTOMER::C_PHONE].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(C_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[CUSTOMER::C_PHONE]);
+ const int32_t *c_nationkey = reinterpret_cast<const int32_t *>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->data());
+ char *c_phone = reinterpret_cast<char *>(
+ tld.batch[CUSTOMER::C_PHONE].array()->buffers[1]->mutable_data());
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ {
+ GeneratePhoneNumber(
+ c_phone + irow * byte_width,
+ tld.rng,
+ c_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_ACCTBAL(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[CUSTOMER::C_ACCTBAL].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_ACCTBAL));
+ Decimal128 *c_acctbal = reinterpret_cast<Decimal128 *>(
+ tld.batch[CUSTOMER::C_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_acctbal[irow] = { dist(tld.rng) };
+ }
+ return Status::OK();
+ }
+
+ Status C_MKTSEGMENT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[CUSTOMER::C_MKTSEGMENT].kind() == Datum::NONE)
+ {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_MKTSEGMENT));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[CUSTOMER::C_MKTSEGMENT]);
+ char *c_mktsegment = reinterpret_cast<char *>(
+ tld.batch[CUSTOMER::C_MKTSEGMENT].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int32_t> dist(0, kNumSegments - 1);
+ for(int64_t irow = 0; irow < tld.to_generate; irow++)
+ {
+ char *out = c_mktsegment + irow * byte_width;
+ int str_idx = dist(tld.rng);
+ std::strncpy(out, Segments[str_idx], byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_COMMENT(size_t thread_index)
+ {
+ ThreadLocalData &tld = thread_local_data_[thread_index];
+ if(tld.batch[CUSTOMER::C_COMMENT].kind() == Datum::NONE)
+ {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[CUSTOMER::C_COMMENT], g_text.GenerateComments(tld.to_generate, 29, 116, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData
+ {
+ random::pcg32_fast rng;
+ int64_t custkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_{0};
+ std::atomic<int64_t> rows_generated_ = { 0 };
+ float scale_factor_{0};
+ int64_t batch_size_{0};
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+ };
+
+ class OrdersGenerator : public TpchTableGenerator
+ {
+ public:
+ OrdersGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen))
+ {
+ }
+
+ Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) override
+ {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ gen_->SetOrdersOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override
+ {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for(size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override
+ {
+ return schema_;
+ }
+
+ private:
+ Status ProduceCallback(size_t)
+ {
+ if(done_.load())
+ return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextOrdersBatch());
+ if(!maybe_batch.has_value())
+ {
+ int64_t batches_generated = gen_->orders_batches_generated();
+ if(batches_generated == batches_outputted_.load())
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+ };
+
+ class LineitemGenerator : public TpchTableGenerator
+ {
+ public:
+ LineitemGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen))
+ {}
+
+ Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) override
+ {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ gen_->SetLineItemOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override
+ {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for(size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override
+ {
+ return schema_;
+ }
+
+ private:
+ Status ProduceCallback(size_t)
+ {
+ if(done_.load())
+ return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextLineItemBatch());
+ if(!maybe_batch.has_value())
+ {
+ int64_t batches_generated = gen_->lineitem_batches_generated();
+ if(batches_generated == batches_outputted_.load())
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_([this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+ };
+
+ class NationGenerator : public TpchTableGenerator
+ {
+ public:
+ Status Init(
+ std::vector<std::string> columns,
+ float /*scale_factor*/,
+ int64_t /*batch_size*/) override
+ {
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(
+ columns,
+ kTypes,
+ kNameMap,
+ column_indices_));
+ rng_.seed(arrow_vendored::pcg_extras::seed_seq_from<std::random_device>{});
+ return Status::OK();
+ }
+
+ Status StartProducing(
+ size_t /*num_threads*/,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback /*schedule_task_callback*/) override
+ {
+ std::shared_ptr<Buffer> N_NATIONKEY_buffer = Buffer::Wrap(kNationKey, sizeof(kNationKey));
+ ArrayData N_NATIONKEY_arraydata(int32(), kRowCount, { nullptr, std::move(N_NATIONKEY_buffer) });
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> N_NAME_buffer, AllocateBuffer(kRowCount * kNameByteWidth));
+ char *N_NAME = reinterpret_cast<char *>(N_NAME_buffer->mutable_data());
+ for(size_t i = 0; i < kRowCount; i++)
+ std::strncpy(N_NAME + kNameByteWidth * i, kCountryNames[i], kNameByteWidth);
+ ArrayData N_NAME_arraydata(fixed_size_binary(kNameByteWidth), kRowCount, { nullptr, std::move(N_NAME_buffer) });
+
+ std::shared_ptr<Buffer> N_REGIONKEY_buffer = Buffer::Wrap(kRegionKey, sizeof(kRegionKey));
+ ArrayData N_REGIONKEY_arraydata(int32(), kRowCount, { nullptr, std::move(N_REGIONKEY_buffer) });
+
+ ARROW_ASSIGN_OR_RAISE(Datum N_COMMENT_datum, g_text.GenerateComments(kRowCount, 31, 114, rng_));
+
+ std::vector<Datum> fields =
+ {
+ std::move(N_NATIONKEY_arraydata),
+ std::move(N_NAME_arraydata),
+ std::move(N_REGIONKEY_arraydata),
+ std::move(N_COMMENT_datum)
+ };
+
+ std::vector<Datum> result;
+ for(const int &col : column_indices_)
+ result.push_back(fields[col]);
+ ARROW_ASSIGN_OR_RAISE(ExecBatch batch, ExecBatch::Make(std::move(result)));
+ output_callback(std::move(batch));
+ finished_callback(static_cast<int64_t>(1));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override
+ {
+ return schema_;
+ }
+
+ private:
+ random::pcg32_fast rng_;
+
+ static constexpr size_t kRowCount = 25;
+ static constexpr int32_t kNameByteWidth = 25;
+ const int32_t kNationKey[kRowCount] = { 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24 };
+ const char *kCountryNames[kRowCount] =
+ {
+ "ALGERIA", "ARGENTINA", "BRAZIL",
+ "CANADA", "EGYPT", "ETHIOPIA",
+ "FRANCE", "GERMANY", "INDIA",
+ "INDONESIA", "IRAN", "IRAQ",
+ "JAPAN", "JORDAN", "KENYA",
+ "MOROCCO", "MOZAMBIQUE", "PERU",
+ "CHINA", "ROMANIA", "SAUDI ARABIA",
+ "VIETNAM", "RUSSIA", "UNITED KINGDOM",
+ "UNITED STATES"
+ };
+ const int32_t kRegionKey[kRowCount] = { 0, 1, 1, 1, 4, 0, 3, 3, 2, 2, 4, 4, 2, 4, 0, 0, 0, 1, 2, 3, 4, 2, 3, 3, 1 };
+
+ struct NATION
+ {
+ enum
+ {
+ N_NATIONKEY,
+ N_NAME,
+ N_REGIONKEY,
+ N_COMMENT,
+ };
+ };
+
+ const std::unordered_map<std::string, int> kNameMap =
+ {
+ { "N_NATIONKEY", NATION::N_NATIONKEY },
+ { "N_NAME", NATION::N_NAME },
+ { "N_REGIONKEY", NATION::N_REGIONKEY },
+ { "N_COMMENT", NATION::N_COMMENT },
+ };
+
+ std::vector<std::shared_ptr<DataType>> kTypes =
+ {
+ int32(),
+ fixed_size_binary(kNameByteWidth),
+ int32(),
+ utf8(),
+ };
+
+ std::shared_ptr<Schema> schema_;
+ std::vector<int> column_indices_;
+ };
+
+ class RegionGenerator : public TpchTableGenerator
+ {
+ public:
+ Status Init(
+ std::vector<std::string> columns,
+ float /*scale_factor*/,
+ int64_t /*batch_size*/) override
+ {
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(
+ columns,
+ kTypes,
+ kNameMap,
+ column_indices_));
+ rng_.seed(arrow_vendored::pcg_extras::seed_seq_from<std::random_device>{});
+ return Status::OK();
+ }
+
+ Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback /*schedule_task_callback*/) override
+ {
+ std::shared_ptr<Buffer> R_REGIONKEY_buffer = Buffer::Wrap(kRegionKey, sizeof(kRegionKey));
+ ArrayData R_REGIONKEY_arraydata(int32(), kRowCount, { nullptr, std::move(R_REGIONKEY_buffer) });
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> R_NAME_buffer, AllocateBuffer(kRowCount * kNameByteWidth));
+ char *R_NAME_data = reinterpret_cast<char *>(R_NAME_buffer->mutable_data());
+ for(size_t i = 0; i < kRowCount; i++)
+ std::strncpy(R_NAME_data + kNameByteWidth * i, kRegionNames[i], kNameByteWidth);
+ ArrayData R_NAME_arraydata(kTypes[static_cast<int>(REGION::R_NAME)], kRowCount, { nullptr, std::move(R_NAME_buffer) });
+
+ ARROW_ASSIGN_OR_RAISE(Datum R_COMMENT_datum, g_text.GenerateComments(kRowCount, 31, 115, rng_));
+
+ std::vector<Datum> fields = { std::move(R_REGIONKEY_arraydata), std::move(R_NAME_arraydata), std::move(R_COMMENT_datum) };
+ std::vector<Datum> result;
+ for(const int &col : column_indices_)
+ result.push_back(fields[col]);
+ ARROW_ASSIGN_OR_RAISE(ExecBatch batch, ExecBatch::Make(std::move(result)));
+ output_callback(std::move(batch));
+ finished_callback(static_cast<int64_t>(1));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override
+ {
+ return schema_;
+ }
+
+ random::pcg32_fast rng_;
+
+ static constexpr size_t kRowCount = 5;
+ static constexpr int32_t kNameByteWidth = 25;
+ const int32_t kRegionKey[kRowCount] = { 0, 1, 2, 3, 4 };
+ const char *kRegionNames[kRowCount] =
+ {
+ "AFRICA", "AMERICA", "ASIA", "EUROPE", "MIDDLE EAST"
+ };
+
+ struct REGION
+ {
+ enum
+ {
+ R_REGIONKEY,
+ R_NAME,
+ R_COMMENT,
+ kNumColumns,
+ };
+ };
+
+ const std::unordered_map<std::string, int> kNameMap =
+ {
+ { "R_REGIONKEY", REGION::R_REGIONKEY },
+ { "R_NAME", REGION::R_NAME },
+ { "R_COMMENT", REGION::R_COMMENT },
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kTypes =
+ {
+ int32(),
+ fixed_size_binary(kNameByteWidth),
+ utf8(),
+ };
+
+ std::shared_ptr<Schema> schema_;
+ std::vector<int> column_indices_;
+ };
+
+ class TpchNode : public ExecNode
+ {
+ public:
+ TpchNode(ExecPlan *plan,
+ const char *name,
+ std::unique_ptr<TpchTableGenerator> generator)
+ : ExecNode(plan, {}, {}, generator->schema(), /*num_outputs=*/1),
+ name_(name),
+ generator_(std::move(generator))
+ {
+ }
+
+ const char *kind_name() const override
+ {
+ return name_;
+ }
+
+ [[noreturn]]
+ static void NoInputs()
+ {
+ Unreachable("TPC-H node should never have any inputs");
+ }
+
+ [[noreturn]]
+ void InputReceived(ExecNode *, ExecBatch) override
+ {
+ NoInputs();
+ }
+
+ [[noreturn]]
+ void ErrorReceived(ExecNode *, Status) override
+ {
+ NoInputs();
+ }
+
+ [[noreturn]]
+ void InputFinished(ExecNode *, int) override
+ {
+ NoInputs();
+ }
+
+ Status StartProducing() override
+ {
+ finished_ = Future<>::Make();
Review comment:
Minor nit: There's no need to do this. Arguably, the base class shouldn't contain `finished_` while leaving `finished()` pure virtual.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829341233
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
Review comment:
switched it to ValidateFull.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829486031
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
Review comment:
Does heap allocation matter for objects with static storage duration? (this is not a rhetorical question, I honestly don't know the answer :) )
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829318988
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<float>(scale_factor));
+
+ ExecNode* lineitem =
+ *gen.Lineitem({"L_QUANTITY", "L_EXTENDEDPRICE", "L_TAX", "L_DISCOUNT", "L_SHIPDATE",
+ "L_RETURNFLAG", "L_LINESTATUS"});
+
+ std::shared_ptr<Date32Scalar> sept_2_1998 = std::make_shared<Date32Scalar>(
+ 10471); // September 2, 1998 is 10471 days after January 1, 1970
+ Expression filter =
+ less_equal(field_ref("L_SHIPDATE"), literal(std::move(sept_2_1998)));
+ FilterNodeOptions filter_opts(filter);
+
+ Expression l_returnflag = field_ref("L_RETURNFLAG");
+ Expression l_linestatus = field_ref("L_LINESTATUS");
+ Expression quantity = field_ref("L_QUANTITY");
+ Expression base_price = field_ref("L_EXTENDEDPRICE");
+
+ std::shared_ptr<Decimal128Scalar> decimal_1 =
+ std::make_shared<Decimal128Scalar>(Decimal128{0, 100}, decimal(12, 2));
+ Expression discount_multiplier =
+ call("subtract", {literal(decimal_1), field_ref("L_DISCOUNT")});
+ Expression tax_multiplier = call("add", {literal(decimal_1), field_ref("L_TAX")});
+ Expression disc_price =
+ call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier});
+ Expression charge =
+ call("multiply",
+ {call("cast",
+ {call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier})},
+ compute::CastOptions::Unsafe(decimal(12, 2))),
+ tax_multiplier});
+ Expression discount = field_ref("L_DISCOUNT");
+
+ std::vector<Expression> projection_list = {l_returnflag, l_linestatus, quantity,
+ base_price, disc_price, charge,
+ quantity, base_price, discount};
+ std::vector<std::string> project_names = {
+ "l_returnflag", "l_linestatus", "sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price", "avg_disc"};
+ ProjectNodeOptions project_opts(std::move(projection_list));
+
+ ScalarAggregateOptions sum_opts = ScalarAggregateOptions::Defaults();
+ CountOptions count_opts(CountOptions::CountMode::ALL);
+ std::vector<arrow::compute::internal::Aggregate> aggs = {
+ {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts},
+ {"hash_sum", &sum_opts}, {"hash_mean", &sum_opts}, {"hash_mean", &sum_opts},
+ {"hash_mean", &sum_opts}, {"hash_count", &count_opts}};
+
+ std::vector<FieldRef> cols = {2, 3, 4, 5, 6, 7, 8, 2};
+
+ std::vector<std::string> names = {"sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price",
+ "avg_disc", "count_order"};
+
+ std::vector<FieldRef> keys = {"L_RETURNFLAG", "L_LINESTATUS"};
+ AggregateNodeOptions agg_opts(aggs, cols, names, keys);
+
+ SortKey l_returnflag_key("L_RETURNFLAG");
+ SortKey l_linestatus_key("L_LINESTATUS");
+ SortOptions sort_opts({l_returnflag_key, l_linestatus_key});
+ OrderBySinkNodeOptions order_by_opts(sort_opts, sink_gen);
+
+ Declaration filter_decl("filter", {Declaration::Input(lineitem)}, filter_opts);
+ Declaration project_decl("project", project_opts);
+ Declaration aggregate_decl("aggregate", agg_opts);
+ Declaration orderby_decl("order_by_sink", order_by_opts);
+
+ Declaration q1 =
+ Declaration::Sequence({filter_decl, project_decl, aggregate_decl, orderby_decl});
+ std::ignore = *q1.AddToPlan(plan.get());
+ return plan;
+}
+
+static void BM_Tpch_Q1(benchmark::State& st) {
+ for (auto _ : st) {
+ st.PauseTiming();
+ AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+ std::shared_ptr<ExecPlan> plan = Plan_Q1(&sink_gen, static_cast<int>(st.range(0)));
+ st.ResumeTiming();
+ auto fut = StartAndCollect(plan.get(), sink_gen);
+ auto res = *fut.MoveResult();
+ }
+}
+
+BENCHMARK(BM_Tpch_Q1)->Args({1})->ArgNames({"SF"});
Review comment:
i think if there's a space, google benchmark will add quotes around it. I could switch it to `ScaleFactor` as a compromise.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r828939116
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3758 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ random::pcg32_fast rng_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ int scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_generated_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
Review comment:
@save-buffer looks like this comment is still unaddressed. Also, if speed is a concern, I notice [~2x faster performance](https://github.com/westonpace/arrow/commit/d9efc98634ceabc09cf187c4d511174efb0eb272) with [`arrow::internal::detail::FormatAllDigitsLeftPadded`](https://github.com/apache/arrow/blob/91f65857c41d580dcf6db563bb2ca776b3eef0d7/cpp/src/arrow/util/formatting.h#L129) than with `AppendNumberPaddedToNineDigits`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r818093649
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,69 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
Review comment:
TIL about that file!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r818743722
##########
File path: r/src/compute-exec.cpp
##########
@@ -277,4 +287,135 @@ std::shared_ptr<compute::ExecNode> ExecNode_ReadFromRecordBatchReader(
return MakeExecNodeOrStop("source", plan.get(), {}, options);
}
+// [[arrow::export]]
+std::shared_ptr<arrow::RecordBatchReader> Tpch_Dbgen(
+ const std::shared_ptr<compute::ExecPlan>& plan,
+ int scale_factor,
+ std::string table_name
+ ) {
+
+ auto gen = ValueOrStop(arrow::compute::TpchGen::Make(plan.get(), scale_factor));
+
+ compute::ExecNode *table;
+ if (table_name == "part") {
+ table = ValueOrStop(gen.Part());
+ } else if (table_name == "supplier") {
+ table = ValueOrStop(gen.Supplier());
+ } else if (table_name == "partsupp") {
+ table = ValueOrStop(gen.PartSupp());
+ } else if (table_name == "customer") {
+ table = ValueOrStop(gen.Customer());
+ } else if (table_name == "nation") {
+ table = ValueOrStop(gen.Nation());
+ } else if (table_name == "lineitem") {
+ table = ValueOrStop(gen.Lineitem());
+ } else if (table_name == "region") {
+ table = ValueOrStop(gen.Region());
+ } else if (table_name == "orders") {
+ table = ValueOrStop(gen.Orders());
+ } else {
+ cpp11::stop("That's not a valid table name");
+ }
+
+ arrow::AsyncGenerator<arrow::util::optional<compute::ExecBatch>> sink_gen;
+
+ MakeExecNodeOrStop("sink", plan.get(), {table},
+ compute::SinkNodeOptions{&sink_gen});
+
+ StopIfNotOk(plan->Validate());
+ StopIfNotOk(plan->StartProducing());
+
+ // If the generator is destroyed before being completely drained, inform plan
+ std::shared_ptr<void> stop_producing{nullptr, [plan](...) {
+ bool not_finished_yet =
+ plan->finished().TryAddCallback([&plan] {
+ return [plan](const arrow::Status&) {};
+ });
+
+ if (not_finished_yet) {
+ plan->StopProducing();
+ }
+ }};
+
+ return compute::MakeGeneratorReader(
+ table->output_schema(),
+ [stop_producing, plan, sink_gen] { return sink_gen(); }, gc_memory_pool());
+}
+
+// [[arrow::export]]
+void Tpch_Dbgen_Write(
Review comment:
This is not the ultimate shape that this function will take — but was my first attempt to using the write node. It is currently segfaulting and I've commented some of the silly things I've done in-line.
This function is a bit separate from the PR so if getting this working will delay merging the larger PR, I'm happy to pull it into a separate one.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r818830313
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3758 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ random::pcg32_fast rng_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ int scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_generated_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
Review comment:
We have formatting utilities in `arrow/util/formatting.h`, no need to reinvent them.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r825008250
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string> &columns,
+ const std::vector<std::shared_ptr<DataType>> &types,
+ const std::unordered_map<std::string, int> &name_map,
+ std::vector<int> &gen_list)
+ {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if(columns.empty())
+ {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for(auto pair : name_map)
+ {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ }
+ else
+ {
+ for(const std::string &col : columns)
+ {
+ auto entry = name_map.find(col);
+ if(entry == name_map.end())
+ return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+ }
+
+ static TpchText g_text;
+
+ Status TpchText::InitIfNeeded(random::pcg32_fast &rng)
+ {
+ if(done_.load())
+ return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(!text_)
+ {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char *out = reinterpret_cast<char *>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+ while(done_.load() == false)
+ {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while(GenerateSentence(try_offset, rng, temp_buff))
+ known_valid_offset = try_offset;
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(done_.load())
+ return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ int64_t memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ std::memcpy(out + generated_offset_, temp_buff, memcpy_size);
+ generated_offset_ += memcpy_size;
+ if(generated_offset_ == kTextBytes)
+ done_.store(true);
+ }
+ }
+ return Status::OK();
+ }
+
+ Result<Datum> TpchText::GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng)
+ {
+ RETURN_NOT_OK(InitIfNeeded(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer, AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for(size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + length_dist(rng);
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer, AllocateBuffer(offsets[num_comments]));
+ char *comments = reinterpret_cast<char *>(comment_buffer->mutable_data());
+ for(size_t i = 0; i < num_comments; i++)
+ {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments, { nullptr, std::move(offset_buffer), std::move(comment_buffer) });
+ return std::move(ad);
+ }
+
+ Result<Datum> RandomVString(
+ random::pcg32_fast &rng,
+ int64_t num_rows,
+ int32_t min_length,
+ int32_t max_length)
+ {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff, AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for(int64_t i = 1; i <= num_rows; i++)
+ offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff, AllocateBuffer(offsets[num_rows]));
+ char *str = reinterpret_cast<char *>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for(int32_t i = 0; i < offsets[num_rows]; i++)
+ str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, { nullptr, std::move(offset_buff), std::move(str_buff) });
+ return std::move(ad);
+ }
+
+ void AppendNumber(char *&out, int num_digits, int32_t x)
+ {
+ out += (num_digits - 1);
+ while(x > 0)
+ {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+ }
+
+ void GeneratePhoneNumber(
+ char *out,
+ random::pcg32_fast &rng,
+ int32_t country)
+ {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+ }
+
+ static constexpr uint32_t STARTDATE = 8035; // January 1, 1992 is 8035 days after January 1, 1970
+ static constexpr uint32_t CURRENTDATE = 9298; // June 17, 1995 is 9298 days after January 1, 1970
+ static constexpr uint32_t ENDDATE = 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+ const char *NameParts[] =
+ {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque", "black", "blanched", "blue",
+ "blush", "brown", "burlywood", "burnished", "chartreuse", "chiffon", "chocolate", "coral",
+ "cornflower", "cornsilk", "cream", "cyan", "dark", "deep", "dim", "dodger", "drab", "firebrick",
+ "floral", "forest", "frosted", "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender", "lawn", "lemon", "light", "lime", "linen",
+ "magenta", "maroon", "medium", "metallic", "midnight", "mint", "misty", "moccasin", "navajo",
+ "navy", "olive", "orange", "orchid", "pale", "papaya", "peach", "peru", "pink", "plum", "powder",
+ "puff", "purple", "red", "rose", "rosy", "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel", "tan", "thistle", "tomato", "turquoise", "violet",
+ "wheat", "white", "yellow",
+ };
+ static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+ const char *Types_1[] =
+ {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+ };
+ static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+ const char *Types_2[] =
+ {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+ };
+ static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+ const char *Types_3[] =
+ {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+ };
+ static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+ const char *Containers_1[] =
+ {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+ };
+ static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+ const char *Containers_2[] =
+ {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+ };
+ static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+ const char *Segments[] =
+ {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+ };
+ static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+ const char *Priorities[] =
+ {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+ };
+ static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+ const char *Instructions[] =
+ {
+ "DELIVER IN PERSON", "COLLECT COD", "NONE", "TAKE BACK RETURN",
+ };
+ static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+ const char *Modes[] =
+ {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+ };
+ static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+ const char *Nouns[] =
+ {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ", "dependencies ", "excuses ",
+ "platelets ", "asymptotes ", "courts ", "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ", "forges ", "braids ", "hockey players ", "frays ",
+ "warhorses ", "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ", "waters ", "orbits ", "gifts ", "sheaves ",
+ "depths ", "sentiments ", "decoys ", "realms ", "pains ", "grouches ", "escapades ",
+ };
+ static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+ const char *Verbs[] =
+ {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ", "boost ", "affix ", "detect ", "integrate ",
+ "maintain ", "nod ", "was ", "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ", "poach ", "serve ", "run ", "dazzle ",
+ "snooze ", "doze ", "unwind ", "kindle ", "play ", "hang ", "believe ", "doubt ",
+ };
+ static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+ const char *Adjectives[] =
+ {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ", "quiet ", "ruthless ", "thin ",
+ "close ", "dogged ", "daring ", "brave ", "stealthy ", "permanent ", "enticing ", "idle ", "busy ",
+ "regular ", "final ", "ironic ", "even ", "bold ", "silent ",
+ };
+ static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+ const char *Adverbs[] =
+ {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ", "blithely ", "quickly ", "fluffily ",
+ "slowly ", "quietly ", "ruthlessly ", "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ", "ironically ", "evenly ", "boldly ",
+ "silently ",
+ };
+ static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+ const char *Prepositions[] =
+ {
+ "about ", "above ", "according to ", "across ", "after ", "against ", "along ", "alongside of ", "among ",
+ "around ", "at ", "atop ", "before ", "behind ", "beneath ", "beside ", "besides ", "between ", "beyond ",
+ "beyond ", "by ", "despite ", "during ", "except ", "for ", "from ", "in place of ", "inside ", "instead of ",
+ "into ", "near ", "of ", "on ", "outside ", "over ", "past ", "since ", "through ", "throughout ", "to ",
+ "toward ", "under ", "until ", "up ", "upon ", "without ", "with ", "within ",
+ };
+ static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+ const char *Auxiliaries[] =
+ {
+ "do ", "may ", "might ", "shall ", "will ", "would ", "can ", "could ", "should ", "ought to ", "must ",
+ "will have to ", "shall have to ", "could have to ", "should have to ", "must have to ", "need to ", "try to ",
+ };
+ static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+ const char *Terminators[] =
+ {
+ ".", ";", ":", "?", "!", "--",
+ };
+ static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+ bool TpchText::GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices)
+ {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char *word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if(offset + length > kChunkSize)
+ return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+ }
+
+ bool TpchText::GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+ }
+
+ bool TpchText::GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+ }
+
+ bool TpchText::GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+ }
+
+ bool TpchText::GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+ }
+
+ bool TpchText::GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+ }
+
+ bool TpchText::GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+ }
+
+ bool TpchText::GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if(result)
+ std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+ }
+
+ bool TpchText::GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char *comma_space = ", ";
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ GenerateAdverb(offset, rng, arr);
+ GenerateAdjective(offset, rng, arr);
+ GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ bool TpchText::GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ bool TpchText::GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ const char *the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+ }
+
+ bool TpchText::GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 5:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ using GenerateColumnFn = std::function<Status(size_t)>;
+ class PartAndPartSupplierGenerator
+ {
+ public:
+ Status Init(
+ size_t num_threads,
+ int64_t batch_size,
+ float scale_factor)
+ {
+ if(!inited_)
+ {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ thread_local_data_.resize(num_threads);
+ for(ThreadLocalData &tld : thread_local_data_)
+ {
+ // 5 is the maximum number of different strings we need to concatenate
+ tld.string_indices.resize(5 * batch_size_);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const
+ {
+ return part_batches_generated_.load();
+ }
+
+ int64_t partsupp_batches_generated() const
+ {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, part_types_, part_name_map_, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, partsupp_types_, partsupp_name_map_, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch()
+ {
+ size_t thread_index = thread_indexer_();
Review comment:
This is one of the great travesties of our time. Since the two ExecNodes have separate thread indexers, different threads may be mapped to the same index, which is a big problem. Ideally the thread index would be a thread_local variable, so that we don't have to grab a mutex to a hash table just to know which index our thread has. The technology just isn't there yet it seems.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r824998653
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string> &columns,
+ const std::vector<std::shared_ptr<DataType>> &types,
+ const std::unordered_map<std::string, int> &name_map,
+ std::vector<int> &gen_list)
+ {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if(columns.empty())
+ {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for(auto pair : name_map)
+ {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ }
+ else
+ {
+ for(const std::string &col : columns)
+ {
+ auto entry = name_map.find(col);
+ if(entry == name_map.end())
+ return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+ }
+
+ static TpchText g_text;
+
+ Status TpchText::InitIfNeeded(random::pcg32_fast &rng)
+ {
+ if(done_.load())
+ return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(!text_)
+ {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char *out = reinterpret_cast<char *>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+ while(done_.load() == false)
+ {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while(GenerateSentence(try_offset, rng, temp_buff))
+ known_valid_offset = try_offset;
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(done_.load())
+ return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ int64_t memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ std::memcpy(out + generated_offset_, temp_buff, memcpy_size);
+ generated_offset_ += memcpy_size;
+ if(generated_offset_ == kTextBytes)
+ done_.store(true);
+ }
Review comment:
I've actually switched it to a lock-free implementation (with an atomic cursor), it seems to have sped up generation by quite a bit.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r825239388
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string> &columns,
+ const std::vector<std::shared_ptr<DataType>> &types,
+ const std::unordered_map<std::string, int> &name_map,
+ std::vector<int> &gen_list)
+ {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if(columns.empty())
+ {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for(auto pair : name_map)
+ {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ }
+ else
+ {
+ for(const std::string &col : columns)
+ {
+ auto entry = name_map.find(col);
+ if(entry == name_map.end())
+ return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+ }
+
+ static TpchText g_text;
+
+ Status TpchText::InitIfNeeded(random::pcg32_fast &rng)
+ {
+ if(done_.load())
+ return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(!text_)
+ {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char *out = reinterpret_cast<char *>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+ while(done_.load() == false)
+ {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while(GenerateSentence(try_offset, rng, temp_buff))
+ known_valid_offset = try_offset;
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(done_.load())
+ return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ int64_t memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ std::memcpy(out + generated_offset_, temp_buff, memcpy_size);
+ generated_offset_ += memcpy_size;
+ if(generated_offset_ == kTextBytes)
+ done_.store(true);
+ }
+ }
+ return Status::OK();
+ }
+
+ Result<Datum> TpchText::GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng)
+ {
+ RETURN_NOT_OK(InitIfNeeded(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer, AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for(size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + length_dist(rng);
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer, AllocateBuffer(offsets[num_comments]));
+ char *comments = reinterpret_cast<char *>(comment_buffer->mutable_data());
+ for(size_t i = 0; i < num_comments; i++)
+ {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments, { nullptr, std::move(offset_buffer), std::move(comment_buffer) });
+ return std::move(ad);
+ }
+
+ Result<Datum> RandomVString(
+ random::pcg32_fast &rng,
+ int64_t num_rows,
+ int32_t min_length,
+ int32_t max_length)
+ {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff, AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for(int64_t i = 1; i <= num_rows; i++)
+ offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff, AllocateBuffer(offsets[num_rows]));
+ char *str = reinterpret_cast<char *>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for(int32_t i = 0; i < offsets[num_rows]; i++)
+ str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, { nullptr, std::move(offset_buff), std::move(str_buff) });
+ return std::move(ad);
+ }
+
+ void AppendNumber(char *&out, int num_digits, int32_t x)
+ {
+ out += (num_digits - 1);
+ while(x > 0)
+ {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+ }
+
+ void GeneratePhoneNumber(
+ char *out,
+ random::pcg32_fast &rng,
+ int32_t country)
+ {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+ }
+
+ static constexpr uint32_t STARTDATE = 8035; // January 1, 1992 is 8035 days after January 1, 1970
+ static constexpr uint32_t CURRENTDATE = 9298; // June 17, 1995 is 9298 days after January 1, 1970
+ static constexpr uint32_t ENDDATE = 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+ const char *NameParts[] =
+ {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque", "black", "blanched", "blue",
+ "blush", "brown", "burlywood", "burnished", "chartreuse", "chiffon", "chocolate", "coral",
+ "cornflower", "cornsilk", "cream", "cyan", "dark", "deep", "dim", "dodger", "drab", "firebrick",
+ "floral", "forest", "frosted", "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender", "lawn", "lemon", "light", "lime", "linen",
+ "magenta", "maroon", "medium", "metallic", "midnight", "mint", "misty", "moccasin", "navajo",
+ "navy", "olive", "orange", "orchid", "pale", "papaya", "peach", "peru", "pink", "plum", "powder",
+ "puff", "purple", "red", "rose", "rosy", "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel", "tan", "thistle", "tomato", "turquoise", "violet",
+ "wheat", "white", "yellow",
+ };
+ static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+ const char *Types_1[] =
+ {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+ };
+ static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+ const char *Types_2[] =
+ {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+ };
+ static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+ const char *Types_3[] =
+ {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+ };
+ static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+ const char *Containers_1[] =
+ {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+ };
+ static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+ const char *Containers_2[] =
+ {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+ };
+ static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+ const char *Segments[] =
+ {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+ };
+ static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+ const char *Priorities[] =
+ {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+ };
+ static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+ const char *Instructions[] =
+ {
+ "DELIVER IN PERSON", "COLLECT COD", "NONE", "TAKE BACK RETURN",
+ };
+ static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+ const char *Modes[] =
+ {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+ };
+ static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+ const char *Nouns[] =
+ {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ", "dependencies ", "excuses ",
+ "platelets ", "asymptotes ", "courts ", "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ", "forges ", "braids ", "hockey players ", "frays ",
+ "warhorses ", "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ", "waters ", "orbits ", "gifts ", "sheaves ",
+ "depths ", "sentiments ", "decoys ", "realms ", "pains ", "grouches ", "escapades ",
+ };
+ static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+ const char *Verbs[] =
+ {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ", "boost ", "affix ", "detect ", "integrate ",
+ "maintain ", "nod ", "was ", "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ", "poach ", "serve ", "run ", "dazzle ",
+ "snooze ", "doze ", "unwind ", "kindle ", "play ", "hang ", "believe ", "doubt ",
+ };
+ static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+ const char *Adjectives[] =
+ {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ", "quiet ", "ruthless ", "thin ",
+ "close ", "dogged ", "daring ", "brave ", "stealthy ", "permanent ", "enticing ", "idle ", "busy ",
+ "regular ", "final ", "ironic ", "even ", "bold ", "silent ",
+ };
+ static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+ const char *Adverbs[] =
+ {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ", "blithely ", "quickly ", "fluffily ",
+ "slowly ", "quietly ", "ruthlessly ", "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ", "ironically ", "evenly ", "boldly ",
+ "silently ",
+ };
+ static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+ const char *Prepositions[] =
+ {
+ "about ", "above ", "according to ", "across ", "after ", "against ", "along ", "alongside of ", "among ",
+ "around ", "at ", "atop ", "before ", "behind ", "beneath ", "beside ", "besides ", "between ", "beyond ",
+ "beyond ", "by ", "despite ", "during ", "except ", "for ", "from ", "in place of ", "inside ", "instead of ",
+ "into ", "near ", "of ", "on ", "outside ", "over ", "past ", "since ", "through ", "throughout ", "to ",
+ "toward ", "under ", "until ", "up ", "upon ", "without ", "with ", "within ",
+ };
+ static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+ const char *Auxiliaries[] =
+ {
+ "do ", "may ", "might ", "shall ", "will ", "would ", "can ", "could ", "should ", "ought to ", "must ",
+ "will have to ", "shall have to ", "could have to ", "should have to ", "must have to ", "need to ", "try to ",
+ };
+ static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+ const char *Terminators[] =
+ {
+ ".", ";", ":", "?", "!", "--",
+ };
+ static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+ bool TpchText::GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices)
+ {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char *word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if(offset + length > kChunkSize)
+ return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+ }
+
+ bool TpchText::GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+ }
+
+ bool TpchText::GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+ }
+
+ bool TpchText::GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+ }
+
+ bool TpchText::GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+ }
+
+ bool TpchText::GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+ }
+
+ bool TpchText::GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+ }
+
+ bool TpchText::GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if(result)
+ std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+ }
+
+ bool TpchText::GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char *comma_space = ", ";
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ GenerateAdverb(offset, rng, arr);
+ GenerateAdjective(offset, rng, arr);
+ GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ bool TpchText::GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ bool TpchText::GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ const char *the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+ }
+
+ bool TpchText::GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 5:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ using GenerateColumnFn = std::function<Status(size_t)>;
+ class PartAndPartSupplierGenerator
+ {
+ public:
+ Status Init(
+ size_t num_threads,
+ int64_t batch_size,
+ float scale_factor)
+ {
+ if(!inited_)
+ {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ thread_local_data_.resize(num_threads);
+ for(ThreadLocalData &tld : thread_local_data_)
+ {
+ // 5 is the maximum number of different strings we need to concatenate
+ tld.string_indices.resize(5 * batch_size_);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const
+ {
+ return part_batches_generated_.load();
+ }
+
+ int64_t partsupp_batches_generated() const
+ {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, part_types_, part_name_map_, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, partsupp_types_, partsupp_name_map_, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch()
+ {
+ size_t thread_index = thread_indexer_();
Review comment:
> This is one of the great travesties of our time.
Then we should fix it. Once #12468 is merged we can assert that all work is done within the executor for the plan. Then we can do something like: https://github.com/apache/arrow/commit/e2324141e57f8c2a3ed7b0da19e2d4103c64df7c
Would that work?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r824981649
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
Review comment:
D:
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1069619675
Let's go ahead and remove the R stuff from this PR and do it separately. I've removed the R stuff, and it should be as easy as reverse-applying this patch:
https://github.com/apache/arrow/commit/b9812de4f4c2ebd1156ac06c29adba37c6d1bade.patch
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1057564538
Yeah, I'm pretty sure a single plan could have one tpch gen node for each table and then each of those nodes would be attached to a dedicated write node. I have done very little testing with multiple sink nodes though so it would be an interesting experiment.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1077337983
I get a sporadic crash when running the benchmarks:
```console
$ ./build-release/relwithdebinfo/arrow-compute-tpch-benchmark
2022-03-24T08:53:30+01:00
Running ./build-release/relwithdebinfo/arrow-compute-tpch-benchmark
Run on (24 X 3800 MHz CPU s)
CPU Caches:
L1 Data 32 KiB (x12)
L1 Instruction 32 KiB (x12)
L2 Unified 512 KiB (x12)
L3 Unified 16384 KiB (x4)
Load Average: 7.97, 7.82, 4.27
***WARNING*** CPU scaling is enabled, the benchmark real time measurements may be noisy and will incur extra overhead.
/home/antoine/arrow/dev/cpp/src/arrow/result.cc:28: ValueOrDie called on an error: Cancelled: Ignoring task added after the task group has been ended
/home/antoine/miniconda3/envs/pyarrow/lib/libarrow.so.800(+0x59505a)[0x7f2d0027605a]
/home/antoine/miniconda3/envs/pyarrow/lib/libarrow.so.800(+0x594ff8)[0x7f2d00275ff8]
/home/antoine/miniconda3/envs/pyarrow/lib/libarrow.so.800(+0x595019)[0x7f2d00276019]
/home/antoine/miniconda3/envs/pyarrow/lib/libarrow.so.800(_ZN5arrow4util8ArrowLogD1Ev+0x21)[0x7f2d00275f31]
/home/antoine/miniconda3/envs/pyarrow/lib/libarrow.so.800(_ZN5arrow8internal17InvalidValueOrDieERKNS_6StatusE+0x175)[0x7f2d00105c55]
./build-release/relwithdebinfo/arrow-compute-tpch-benchmark[0x40fb69]
/home/antoine/miniconda3/envs/pyarrow/lib/libbenchmark.so.1(_ZNK9benchmark8internal17BenchmarkInstance3RunEmiPNS0_11ThreadTimerEPNS0_13ThreadManagerEPNS0_23PerfCountersMeasurementE+0xa8)[0x7f2d01a724c8]
/home/antoine/miniconda3/envs/pyarrow/lib/libbenchmark.so.1(+0x4536d)[0x7f2d01a8f36d]
/home/antoine/miniconda3/envs/pyarrow/lib/libbenchmark.so.1(_ZN9benchmark8internal15BenchmarkRunner13DoNIterationsEv+0x3ad)[0x7f2d01a8fa6d]
/home/antoine/miniconda3/envs/pyarrow/lib/libbenchmark.so.1(_ZN9benchmark8internal15BenchmarkRunner15DoOneRepetitionEv+0x11e)[0x7f2d01a91f9e]
/home/antoine/miniconda3/envs/pyarrow/lib/libbenchmark.so.1(+0x25d84)[0x7f2d01a6fd84]
/home/antoine/miniconda3/envs/pyarrow/lib/libbenchmark.so.1(_ZN9benchmark22RunSpecifiedBenchmarksEPNS_17BenchmarkReporterES1_NSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE+0x582)[0x7f2d01a710e2]
/home/antoine/miniconda3/envs/pyarrow/lib/libbenchmark.so.1(_ZN9benchmark22RunSpecifiedBenchmarksEv+0x7b)[0x7f2d01a712fb]
/home/antoine/miniconda3/envs/pyarrow/lib/libbenchmark_main.so.1(main+0x37)[0x7f2d01abf067]
/lib/x86_64-linux-gnu/libc.so.6(__libc_start_main+0xf3)[0x7f2cff9200b3]
./build-release/relwithdebinfo/arrow-compute-tpch-benchmark[0x40432e]
Abandon (core dumped)
```
Here is the gdb traceback:
```
(gdb) bt
#0 __GI_raise (sig=sig@entry=6) at ../sysdeps/unix/sysv/linux/raise.c:50
#1 0x00007f2cff91e859 in __GI_abort () at abort.c:79
#2 0x00007f2d00275ffe in arrow::util::CerrLog::~CerrLog (this=0x1654e80) at /home/antoine/arrow/dev/cpp/src/arrow/util/logging.cc:72
#3 0x00007f2d00276019 in arrow::util::CerrLog::~CerrLog (this=0x1654e80) at /home/antoine/arrow/dev/cpp/src/arrow/util/logging.cc:66
#4 0x00007f2d00275f31 in arrow::util::ArrowLog::~ArrowLog (this=0x7fffebdf4a50) at /home/antoine/arrow/dev/cpp/src/arrow/util/logging.cc:250
#5 0x00007f2d00105c55 in arrow::internal::DieWithMessage (msg="ValueOrDie called on an error: Cancelled: Ignoring task added after the task group has been ended")
at /home/antoine/arrow/dev/cpp/src/arrow/result.cc:28
#6 arrow::internal::InvalidValueOrDie (st=...) at /home/antoine/arrow/dev/cpp/src/arrow/result.cc:31
#7 0x000000000040fb69 in arrow::Result<std::vector<arrow::compute::ExecBatch, std::allocator<arrow::compute::ExecBatch> > >::ValueOrDie() && (this=0x7f2c980105b0)
at /home/antoine/arrow/dev/cpp/src/arrow/result.h:339
#8 arrow::Result<std::vector<arrow::compute::ExecBatch, std::allocator<arrow::compute::ExecBatch> > >::operator*() && (this=0x7f2c980105b0)
at /home/antoine/arrow/dev/cpp/src/arrow/result.h:343
#9 arrow::compute::internal::BM_Tpch_Q1 (st=...) at /home/antoine/arrow/dev/cpp/src/arrow/compute/exec/tpch_benchmark.cc:116
#10 0x00007f2d01a724c8 in benchmark::internal::BenchmarkInstance::Run(unsigned long, int, benchmark::internal::ThreadTimer*, benchmark::internal::ThreadManager*, benchmark::internal::PerfCountersMeasurement*) const () from /home/antoine/miniconda3/envs/pyarrow/lib/libbenchmark.so.1
#11 0x00007f2d01a8f36d in benchmark::internal::(anonymous namespace)::RunInThread(benchmark::internal::BenchmarkInstance const*, unsigned long, int, benchmark::internal::ThreadManager*, benchmark::internal::PerfCountersMeasurement*) () from /home/antoine/miniconda3/envs/pyarrow/lib/libbenchmark.so.1
#12 0x00007f2d01a8fa6d in benchmark::internal::BenchmarkRunner::DoNIterations() () from /home/antoine/miniconda3/envs/pyarrow/lib/libbenchmark.so.1
#13 0x00007f2d01a91f9e in benchmark::internal::BenchmarkRunner::DoOneRepetition() () from /home/antoine/miniconda3/envs/pyarrow/lib/libbenchmark.so.1
#14 0x00007f2d01a6fd84 in benchmark::internal::(anonymous namespace)::RunBenchmarks(std::vector<benchmark::internal::BenchmarkInstance, std::allocator<benchmark::internal::BenchmarkInstance> > const&, benchmark::BenchmarkReporter*, benchmark::BenchmarkReporter*) () from /home/antoine/miniconda3/envs/pyarrow/lib/libbenchmark.so.1
#15 0x00007f2d01a710e2 in benchmark::RunSpecifiedBenchmarks(benchmark::BenchmarkReporter*, benchmark::BenchmarkReporter*, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >) () from /home/antoine/miniconda3/envs/pyarrow/lib/libbenchmark.so.1
#16 0x00007f2d01a712fb in benchmark::RunSpecifiedBenchmarks() () from /home/antoine/miniconda3/envs/pyarrow/lib/libbenchmark.so.1
#17 0x00007f2d01abf067 in main () from /home/antoine/miniconda3/envs/pyarrow/lib/libbenchmark_main.so.1
#18 0x00007f2cff9200b3 in __libc_start_main (main=0x7f2d01abf030 <main>, argc=1, argv=0x7fffebdf7f08, init=<optimized out>, fini=<optimized out>,
rtld_fini=<optimized out>, stack_end=0x7fffebdf7ef8) at ../csu/libc-start.c:308
#19 0x000000000040432e in _start () at /home/antoine/arrow/dev/cpp/src/arrow/compute/exec/tpch_benchmark.cc:120
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082232256
> I think either way we don't have to make the factory return a unique_ptr.
Indeed, we don't. We can use the pimpl idiom instead (which might be slightly more verbose). The virtual methods seem harmless to me, given that the function call cost is not critical here, but either way is ok :-)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] ursabot commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
ursabot commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082832095
Benchmark runs are scheduled for baseline = 0deefd1838bf4706232b90b72da010698010f7b6 and contender = 50fab73da58c9f76fbe675d661e56087e2f5fb52. 50fab73da58c9f76fbe675d661e56087e2f5fb52 is a master commit associated with this PR. Results will be available as each benchmark for each run completes.
Conbench compare runs links:
[Scheduled] [ec2-t3-xlarge-us-east-2](https://conbench.ursa.dev/compare/runs/3f8ac59fddef426cac31b6b2826ce5a6...e5b78c70b6fe4a699c4cb8b967fe5736/)
[Scheduled] [test-mac-arm](https://conbench.ursa.dev/compare/runs/02b74ab900d440f8b5361300f15ad15c...b870e78b617e41e98c2ade7241e0aa06/)
[Scheduled] [ursa-i9-9960x](https://conbench.ursa.dev/compare/runs/1f7740546747452aa61fece1295aa0e3...f66d8edd43724d9988628cbe4a4e4d50/)
[Scheduled] [ursa-thinkcentre-m75q](https://conbench.ursa.dev/compare/runs/c3c55d6999fe4d7eabec73cbdfa1fad6...4b9f45f6442447d3a9d758598d0e14d3/)
Supported benchmarks:
ec2-t3-xlarge-us-east-2: Supported benchmark langs: Python, R. Runs only benchmarks with cloud = True
test-mac-arm: Supported benchmark langs: C++, Python, R
ursa-i9-9960x: Supported benchmark langs: Python, R, JavaScript
ursa-thinkcentre-m75q: Supported benchmark langs: C++, Java
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] ursabot edited a comment on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
ursabot edited a comment on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082832095
Benchmark runs are scheduled for baseline = 0deefd1838bf4706232b90b72da010698010f7b6 and contender = 50fab73da58c9f76fbe675d661e56087e2f5fb52. 50fab73da58c9f76fbe675d661e56087e2f5fb52 is a master commit associated with this PR. Results will be available as each benchmark for each run completes.
Conbench compare runs links:
[Finished :arrow_down:0.0% :arrow_up:0.0%] [ec2-t3-xlarge-us-east-2](https://conbench.ursa.dev/compare/runs/3f8ac59fddef426cac31b6b2826ce5a6...e5b78c70b6fe4a699c4cb8b967fe5736/)
[Scheduled] [test-mac-arm](https://conbench.ursa.dev/compare/runs/02b74ab900d440f8b5361300f15ad15c...b870e78b617e41e98c2ade7241e0aa06/)
[Scheduled] [ursa-i9-9960x](https://conbench.ursa.dev/compare/runs/1f7740546747452aa61fece1295aa0e3...f66d8edd43724d9988628cbe4a4e4d50/)
[Scheduled] [ursa-thinkcentre-m75q](https://conbench.ursa.dev/compare/runs/c3c55d6999fe4d7eabec73cbdfa1fad6...4b9f45f6442447d3a9d758598d0e14d3/)
Buildkite builds:
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/415| `50fab73d` ec2-t3-xlarge-us-east-2>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/414| `0deefd18` ec2-t3-xlarge-us-east-2>
[Scheduled] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-test-mac-arm/builds/400| `0deefd18` test-mac-arm>
[Scheduled] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-thinkcentre-m75q/builds/410| `0deefd18` ursa-thinkcentre-m75q>
Supported benchmarks:
ec2-t3-xlarge-us-east-2: Supported benchmark langs: Python, R. Runs only benchmarks with cloud = True
test-mac-arm: Supported benchmark langs: C++, Python, R
ursa-i9-9960x: Supported benchmark langs: Python, R, JavaScript
ursa-thinkcentre-m75q: Supported benchmark langs: C++, Java
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] ursabot edited a comment on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
ursabot edited a comment on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082832095
Benchmark runs are scheduled for baseline = 0deefd1838bf4706232b90b72da010698010f7b6 and contender = 50fab73da58c9f76fbe675d661e56087e2f5fb52. 50fab73da58c9f76fbe675d661e56087e2f5fb52 is a master commit associated with this PR. Results will be available as each benchmark for each run completes.
Conbench compare runs links:
[Finished :arrow_down:0.0% :arrow_up:0.0%] [ec2-t3-xlarge-us-east-2](https://conbench.ursa.dev/compare/runs/3f8ac59fddef426cac31b6b2826ce5a6...e5b78c70b6fe4a699c4cb8b967fe5736/)
[Finished :arrow_down:0.38% :arrow_up:0.33%] [test-mac-arm](https://conbench.ursa.dev/compare/runs/02b74ab900d440f8b5361300f15ad15c...b870e78b617e41e98c2ade7241e0aa06/)
[Finished :arrow_down:2.14% :arrow_up:0.0%] [ursa-i9-9960x](https://conbench.ursa.dev/compare/runs/1f7740546747452aa61fece1295aa0e3...f66d8edd43724d9988628cbe4a4e4d50/)
[Finished :arrow_down:0.04% :arrow_up:0.26%] [ursa-thinkcentre-m75q](https://conbench.ursa.dev/compare/runs/c3c55d6999fe4d7eabec73cbdfa1fad6...4b9f45f6442447d3a9d758598d0e14d3/)
Buildkite builds:
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/415| `50fab73d` ec2-t3-xlarge-us-east-2>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-test-mac-arm/builds/401| `50fab73d` test-mac-arm>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-i9-9960x/builds/401| `50fab73d` ursa-i9-9960x>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-thinkcentre-m75q/builds/411| `50fab73d` ursa-thinkcentre-m75q>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/414| `0deefd18` ec2-t3-xlarge-us-east-2>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-test-mac-arm/builds/400| `0deefd18` test-mac-arm>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-i9-9960x/builds/400| `0deefd18` ursa-i9-9960x>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-thinkcentre-m75q/builds/410| `0deefd18` ursa-thinkcentre-m75q>
Supported benchmarks:
ec2-t3-xlarge-us-east-2: Supported benchmark langs: Python, R. Runs only benchmarks with cloud = True
test-mac-arm: Supported benchmark langs: C++, Python, R
ursa-i9-9960x: Supported benchmark langs: Python, R, JavaScript
ursa-thinkcentre-m75q: Supported benchmark langs: C++, Java
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] github-actions[bot] commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1055989579
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r817330103
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3738 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status Init();
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ void GenerateWord(size_t &offset, const char **words, size_t num_choices);
+ void GenerateNoun(size_t &offset);
+ void GenerateVerb(size_t &offset);
+ void GenerateAdjective(size_t &offset);
+ void GenerateAdverb(size_t &offset);
+ void GeneratePreposition(size_t &offset);
+ void GenerateAuxiliary(size_t &offset);
+ void GenerateTerminator(size_t &offset);
+
+ void GenerateNounPhrase(size_t &offset);
+ void GenerateVerbPhrase(size_t &offset);
+ void GeneratePrepositionalPhrase(size_t &offset);
+
+ void GenerateSentence(size_t &offset);
+
+ std::unique_ptr<Buffer> text_;
+ random::pcg32_fast rng_;
+ static constexpr size_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ int scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_generated_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
Review comment:
Modification doesn't really matter from a legal perspective. Derivative works are only allowed if you have rights to the source.
That being said, we are in a terrible gray area here that I fear will become a can of worms and a waste of everyone's time. Technically the only "proper" approach would be a [clean room approach](https://en.wikipedia.org/wiki/Clean_room_design) where @save-buffer describes what is needed to someone that has never seen the SO code and that person writes the code. However, legally documenting such a process is a headache.
In this case I think we are (at least ethically) in the clear. The answer author has this statement on their profile page:
> All code I post on Stack Overflow is covered by the "Do whatever the heck you want with it" licence, the full text of which is:
>
> Do whatever the heck you want with it.
>
There are a few edits by other authors but none of which touches the code included here. My opinion would be to proceed as is (but get rid of the ASCII art, I can't abide fun).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829349153
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
Review comment:
yes, the spec specifies to generate in the range 100-999 and 1000-9999
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829563018
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
Review comment:
As an end user of this code, I would appreciate the ability to provide determinism if that's not too painful.
We wouldn't always want or need it, but sometimes it is helpful to be able to generate the exact same set of data.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829466609
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
+class OrdersAndLineItemGenerator;
+class PartAndPartSupplierGenerator;
+
+class ARROW_EXPORT TpchGen {
+ public:
+ /*
+ * \brief Create a factory for nodes that generate TPC-H data
+ *
+ * Note: Individual tables will reference each other. It is important that you only
+ * create a single TpchGen instance for each plan and then you can create nodes for each
+ * table from that single TpchGen instance. Note: Every batch will be scheduled as a new
+ * task using the ExecPlan's scheduler.
+ */
+ static Result<TpchGen> Make(ExecPlan* plan, float scale_factor = 1.0f,
+ int64_t batch_size = 4096);
Review comment:
There's actually nothing preventing us from doing the two batch-size approach today:
```
void OutputBatchCallback(ExecBatch batch) {
for (const auto& small_batch : Partition(batch, small_batch_size)) {
outputs_[0]->InputReceived(this, small_batch);
}
}
```
However:
* The vector/shared_ptr allocations introduced by slicing will fight the benefits
* The filter and project nodes right now are unfortunately creating a thread task for every batch which will sort of defeat the purpose.
I agree this is best left for a follow-up. @sanjibansg is doing some experiments here over the next month or two as part of an internship project.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] ursabot edited a comment on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
ursabot edited a comment on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082832095
Benchmark runs are scheduled for baseline = 0deefd1838bf4706232b90b72da010698010f7b6 and contender = 50fab73da58c9f76fbe675d661e56087e2f5fb52. 50fab73da58c9f76fbe675d661e56087e2f5fb52 is a master commit associated with this PR. Results will be available as each benchmark for each run completes.
Conbench compare runs links:
[Finished :arrow_down:0.0% :arrow_up:0.0%] [ec2-t3-xlarge-us-east-2](https://conbench.ursa.dev/compare/runs/3f8ac59fddef426cac31b6b2826ce5a6...e5b78c70b6fe4a699c4cb8b967fe5736/)
[Scheduled] [test-mac-arm](https://conbench.ursa.dev/compare/runs/02b74ab900d440f8b5361300f15ad15c...b870e78b617e41e98c2ade7241e0aa06/)
[Scheduled] [ursa-i9-9960x](https://conbench.ursa.dev/compare/runs/1f7740546747452aa61fece1295aa0e3...f66d8edd43724d9988628cbe4a4e4d50/)
[Finished :arrow_down:0.04% :arrow_up:0.26%] [ursa-thinkcentre-m75q](https://conbench.ursa.dev/compare/runs/c3c55d6999fe4d7eabec73cbdfa1fad6...4b9f45f6442447d3a9d758598d0e14d3/)
Buildkite builds:
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/415| `50fab73d` ec2-t3-xlarge-us-east-2>
[Scheduled] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-test-mac-arm/builds/401| `50fab73d` test-mac-arm>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-thinkcentre-m75q/builds/411| `50fab73d` ursa-thinkcentre-m75q>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/414| `0deefd18` ec2-t3-xlarge-us-east-2>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-test-mac-arm/builds/400| `0deefd18` test-mac-arm>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-thinkcentre-m75q/builds/410| `0deefd18` ursa-thinkcentre-m75q>
Supported benchmarks:
ec2-t3-xlarge-us-east-2: Supported benchmark langs: Python, R. Runs only benchmarks with cloud = True
test-mac-arm: Supported benchmark langs: C++, Python, R
ursa-i9-9960x: Supported benchmark langs: Python, R, JavaScript
ursa-thinkcentre-m75q: Supported benchmark langs: C++, Java
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] ursabot edited a comment on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
ursabot edited a comment on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082832095
Benchmark runs are scheduled for baseline = 0deefd1838bf4706232b90b72da010698010f7b6 and contender = 50fab73da58c9f76fbe675d661e56087e2f5fb52. 50fab73da58c9f76fbe675d661e56087e2f5fb52 is a master commit associated with this PR. Results will be available as each benchmark for each run completes.
Conbench compare runs links:
[Finished :arrow_down:0.0% :arrow_up:0.0%] [ec2-t3-xlarge-us-east-2](https://conbench.ursa.dev/compare/runs/3f8ac59fddef426cac31b6b2826ce5a6...e5b78c70b6fe4a699c4cb8b967fe5736/)
[Scheduled] [test-mac-arm](https://conbench.ursa.dev/compare/runs/02b74ab900d440f8b5361300f15ad15c...b870e78b617e41e98c2ade7241e0aa06/)
[Scheduled] [ursa-i9-9960x](https://conbench.ursa.dev/compare/runs/1f7740546747452aa61fece1295aa0e3...f66d8edd43724d9988628cbe4a4e4d50/)
[Scheduled] [ursa-thinkcentre-m75q](https://conbench.ursa.dev/compare/runs/c3c55d6999fe4d7eabec73cbdfa1fad6...4b9f45f6442447d3a9d758598d0e14d3/)
Buildkite builds:
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/415| `50fab73d` ec2-t3-xlarge-us-east-2>
[Scheduled] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-thinkcentre-m75q/builds/411| `50fab73d` ursa-thinkcentre-m75q>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/414| `0deefd18` ec2-t3-xlarge-us-east-2>
[Scheduled] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-test-mac-arm/builds/400| `0deefd18` test-mac-arm>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-thinkcentre-m75q/builds/410| `0deefd18` ursa-thinkcentre-m75q>
Supported benchmarks:
ec2-t3-xlarge-us-east-2: Supported benchmark langs: Python, R. Runs only benchmarks with cloud = True
test-mac-arm: Supported benchmark langs: C++, Python, R
ursa-i9-9960x: Supported benchmark langs: Python, R, JavaScript
ursa-thinkcentre-m75q: Supported benchmark langs: C++, Java
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r824980006
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
Review comment:
This is how the hash join node does it, so I copied it.
As a general note, I like the style of passing functions around is _much_ more than these objects because a) it's way more boilerplate code to define these classes and b) semantically it feels like objects should be for bundling state with operations. This type of stuff just feels like a hack to pass a v-table around implicitly.
Also I'd argue this is more Java-vs-functional-style
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829665966
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ if (!partsupp_output_queue_.empty()) {
+ ExecBatch result = std::move(partsupp_output_queue_.front());
+ partsupp_output_queue_.pop();
+ return std::move(result);
+ }
+ }
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+ if (!part_cols_.empty()) {
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch part_batch,
+ ExecBatch::Make(std::move(part_result)));
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ part_output_queue_.emplace(std::move(part_batch));
+ }
+ }
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ for (size_t i = 1; i < partsupp_results.size(); i++)
+ partsupp_output_queue_.emplace(std::move(partsupp_results[i]));
+ }
+ return std::move(partsupp_results[0]);
+ }
+
+ private:
+#define FOR_EACH_PART_COLUMN(F) \
+ F(P_PARTKEY) \
+ F(P_NAME) \
+ F(P_MFGR) \
+ F(P_BRAND) \
+ F(P_TYPE) \
+ F(P_SIZE) \
+ F(P_CONTAINER) \
+ F(P_RETAILPRICE) \
+ F(P_COMMENT)
+
+#define FOR_EACH_PARTSUPP_COLUMN(F) \
+ F(PS_PARTKEY) \
+ F(PS_SUPPKEY) \
+ F(PS_AVAILQTY) \
+ F(PS_SUPPLYCOST) \
+ F(PS_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct PART {
+ enum {
+ FOR_EACH_PART_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct PARTSUPP {
+ enum {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, PART::col},
+ const std::unordered_map<std::string, int> kPartNameMap = {
+ FOR_EACH_PART_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, PARTSUPP::col},
+ const std::unordered_map<std::string, int> kPartsuppNameMap = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kPartGenerators = {FOR_EACH_PART_COLUMN(MAKE_FN_ARRAY)};
+ std::vector<GenerateColumnFn> kPartsuppGenerators = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kPartTypes = {
+ int32(), utf8(), fixed_size_binary(25), fixed_size_binary(10),
+ utf8(), int32(), fixed_size_binary(10), decimal(12, 2),
+ utf8(),
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kPartsuppTypes = {
+ int32(), int32(), int32(), decimal(12, 2), utf8(),
+ };
+
+ Status AllocatePartBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.part[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.part_to_generate * byte_width));
+ ArrayData ad(kPartTypes[column], tld.part_to_generate, {nullptr, std::move(buff)});
+ tld.part[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status P_PARTKEY(size_t thread_index) {
Review comment:
i added a big one at the top
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829466609
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
+class OrdersAndLineItemGenerator;
+class PartAndPartSupplierGenerator;
+
+class ARROW_EXPORT TpchGen {
+ public:
+ /*
+ * \brief Create a factory for nodes that generate TPC-H data
+ *
+ * Note: Individual tables will reference each other. It is important that you only
+ * create a single TpchGen instance for each plan and then you can create nodes for each
+ * table from that single TpchGen instance. Note: Every batch will be scheduled as a new
+ * task using the ExecPlan's scheduler.
+ */
+ static Result<TpchGen> Make(ExecPlan* plan, float scale_factor = 1.0f,
+ int64_t batch_size = 4096);
Review comment:
There's actually nothing preventing us from doing the two batch-size approach today:
```
void OutputBatchCallback(ExecBatch batch) {
for (const auto& small_batch : Partition(batch, small_batch_size)) {
outputs_[0]->InputReceived(this, small_batch);
}
}
```
However:
* The vector/shared_ptr allocations introduced by slicing will fight the benefits
* The filter and project nodes right now are unfortunately creating a thread task for every batch which will sort of defeat the purpose.
I agree this is best left for a follow-up. @sanjibansg is doing some experiments here over the next month or two as part of an internship project.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
Review comment:
I agree we want users to use this but I don't think that means we will need the node in a public namespace. Eventually, we will add function registry and Substrait integrations. At that point we will probably create some kind of "options" object which will be in the public API. At that point the nodes could even be in anonymous namespace, however I think the direct construction is necessary now and will probably still be useful in the future for unit testing. I think `arrow::compute::internal` is reasonable.
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
Review comment:
This is a relatively recent introduction to the style guide (sadly it is not versioned) that we chose not to adopt: https://issues.apache.org/jira/browse/ARROW-14960
We really should add this fact to the C++ development guidelines. I will make a reminder to check it this evening.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
Review comment:
Does heap allocation matter for objects with static storage duration? (this is not a rhetorical question, I honestly don't know the answer :) )
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
Review comment:
Agreed. We can take a seed in the constructor. If not specified I think defaulting to non-deterministic is ok.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1069621081
Thanks @save-buffer I agree, that will be simpler. I'm happy to take a look at `tpch_dbgen_write` too.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r825248624
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string> &columns,
+ const std::vector<std::shared_ptr<DataType>> &types,
+ const std::unordered_map<std::string, int> &name_map,
+ std::vector<int> &gen_list)
+ {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if(columns.empty())
+ {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for(auto pair : name_map)
+ {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ }
+ else
+ {
+ for(const std::string &col : columns)
+ {
+ auto entry = name_map.find(col);
+ if(entry == name_map.end())
+ return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+ }
+
+ static TpchText g_text;
+
+ Status TpchText::InitIfNeeded(random::pcg32_fast &rng)
+ {
+ if(done_.load())
+ return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(!text_)
+ {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char *out = reinterpret_cast<char *>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+ while(done_.load() == false)
+ {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while(GenerateSentence(try_offset, rng, temp_buff))
+ known_valid_offset = try_offset;
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(done_.load())
+ return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ int64_t memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ std::memcpy(out + generated_offset_, temp_buff, memcpy_size);
+ generated_offset_ += memcpy_size;
+ if(generated_offset_ == kTextBytes)
+ done_.store(true);
+ }
+ }
+ return Status::OK();
+ }
+
+ Result<Datum> TpchText::GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng)
+ {
+ RETURN_NOT_OK(InitIfNeeded(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer, AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for(size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + length_dist(rng);
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer, AllocateBuffer(offsets[num_comments]));
+ char *comments = reinterpret_cast<char *>(comment_buffer->mutable_data());
+ for(size_t i = 0; i < num_comments; i++)
+ {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments, { nullptr, std::move(offset_buffer), std::move(comment_buffer) });
+ return std::move(ad);
+ }
+
+ Result<Datum> RandomVString(
+ random::pcg32_fast &rng,
+ int64_t num_rows,
+ int32_t min_length,
+ int32_t max_length)
+ {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff, AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t *offsets = reinterpret_cast<int32_t *>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for(int64_t i = 1; i <= num_rows; i++)
+ offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff, AllocateBuffer(offsets[num_rows]));
+ char *str = reinterpret_cast<char *>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for(int32_t i = 0; i < offsets[num_rows]; i++)
+ str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, { nullptr, std::move(offset_buff), std::move(str_buff) });
+ return std::move(ad);
+ }
+
+ void AppendNumber(char *&out, int num_digits, int32_t x)
+ {
+ out += (num_digits - 1);
+ while(x > 0)
+ {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+ }
+
+ void GeneratePhoneNumber(
+ char *out,
+ random::pcg32_fast &rng,
+ int32_t country)
+ {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+ }
+
+ static constexpr uint32_t STARTDATE = 8035; // January 1, 1992 is 8035 days after January 1, 1970
+ static constexpr uint32_t CURRENTDATE = 9298; // June 17, 1995 is 9298 days after January 1, 1970
+ static constexpr uint32_t ENDDATE = 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+ const char *NameParts[] =
+ {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque", "black", "blanched", "blue",
+ "blush", "brown", "burlywood", "burnished", "chartreuse", "chiffon", "chocolate", "coral",
+ "cornflower", "cornsilk", "cream", "cyan", "dark", "deep", "dim", "dodger", "drab", "firebrick",
+ "floral", "forest", "frosted", "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender", "lawn", "lemon", "light", "lime", "linen",
+ "magenta", "maroon", "medium", "metallic", "midnight", "mint", "misty", "moccasin", "navajo",
+ "navy", "olive", "orange", "orchid", "pale", "papaya", "peach", "peru", "pink", "plum", "powder",
+ "puff", "purple", "red", "rose", "rosy", "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel", "tan", "thistle", "tomato", "turquoise", "violet",
+ "wheat", "white", "yellow",
+ };
+ static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+ const char *Types_1[] =
+ {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+ };
+ static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+ const char *Types_2[] =
+ {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+ };
+ static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+ const char *Types_3[] =
+ {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+ };
+ static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+ const char *Containers_1[] =
+ {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+ };
+ static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+ const char *Containers_2[] =
+ {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+ };
+ static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+ const char *Segments[] =
+ {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+ };
+ static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+ const char *Priorities[] =
+ {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+ };
+ static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+ const char *Instructions[] =
+ {
+ "DELIVER IN PERSON", "COLLECT COD", "NONE", "TAKE BACK RETURN",
+ };
+ static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+ const char *Modes[] =
+ {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+ };
+ static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+ const char *Nouns[] =
+ {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ", "dependencies ", "excuses ",
+ "platelets ", "asymptotes ", "courts ", "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ", "forges ", "braids ", "hockey players ", "frays ",
+ "warhorses ", "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ", "waters ", "orbits ", "gifts ", "sheaves ",
+ "depths ", "sentiments ", "decoys ", "realms ", "pains ", "grouches ", "escapades ",
+ };
+ static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+ const char *Verbs[] =
+ {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ", "boost ", "affix ", "detect ", "integrate ",
+ "maintain ", "nod ", "was ", "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ", "poach ", "serve ", "run ", "dazzle ",
+ "snooze ", "doze ", "unwind ", "kindle ", "play ", "hang ", "believe ", "doubt ",
+ };
+ static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+ const char *Adjectives[] =
+ {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ", "quiet ", "ruthless ", "thin ",
+ "close ", "dogged ", "daring ", "brave ", "stealthy ", "permanent ", "enticing ", "idle ", "busy ",
+ "regular ", "final ", "ironic ", "even ", "bold ", "silent ",
+ };
+ static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+ const char *Adverbs[] =
+ {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ", "blithely ", "quickly ", "fluffily ",
+ "slowly ", "quietly ", "ruthlessly ", "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ", "ironically ", "evenly ", "boldly ",
+ "silently ",
+ };
+ static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+ const char *Prepositions[] =
+ {
+ "about ", "above ", "according to ", "across ", "after ", "against ", "along ", "alongside of ", "among ",
+ "around ", "at ", "atop ", "before ", "behind ", "beneath ", "beside ", "besides ", "between ", "beyond ",
+ "beyond ", "by ", "despite ", "during ", "except ", "for ", "from ", "in place of ", "inside ", "instead of ",
+ "into ", "near ", "of ", "on ", "outside ", "over ", "past ", "since ", "through ", "throughout ", "to ",
+ "toward ", "under ", "until ", "up ", "upon ", "without ", "with ", "within ",
+ };
+ static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+ const char *Auxiliaries[] =
+ {
+ "do ", "may ", "might ", "shall ", "will ", "would ", "can ", "could ", "should ", "ought to ", "must ",
+ "will have to ", "shall have to ", "could have to ", "should have to ", "must have to ", "need to ", "try to ",
+ };
+ static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+ const char *Terminators[] =
+ {
+ ".", ";", ":", "?", "!", "--",
+ };
+ static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+ bool TpchText::GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices)
+ {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char *word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if(offset + length > kChunkSize)
+ return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+ }
+
+ bool TpchText::GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+ }
+
+ bool TpchText::GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+ }
+
+ bool TpchText::GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+ }
+
+ bool TpchText::GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+ }
+
+ bool TpchText::GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+ }
+
+ bool TpchText::GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+ }
+
+ bool TpchText::GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if(result)
+ std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+ }
+
+ bool TpchText::GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char *comma_space = ", ";
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ GenerateAdverb(offset, rng, arr);
+ GenerateAdjective(offset, rng, arr);
+ GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ bool TpchText::GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ bool TpchText::GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ const char *the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+ }
+
+ bool TpchText::GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr)
+ {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch(dist(rng))
+ {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 5:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+ }
+
+ using GenerateColumnFn = std::function<Status(size_t)>;
+ class PartAndPartSupplierGenerator
+ {
+ public:
+ Status Init(
+ size_t num_threads,
+ int64_t batch_size,
+ float scale_factor)
+ {
+ if(!inited_)
+ {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ thread_local_data_.resize(num_threads);
+ for(ThreadLocalData &tld : thread_local_data_)
+ {
+ // 5 is the maximum number of different strings we need to concatenate
+ tld.string_indices.resize(5 * batch_size_);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const
+ {
+ return part_batches_generated_.load();
+ }
+
+ int64_t partsupp_batches_generated() const
+ {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, part_types_, part_name_map_, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(const std::vector<std::string> &cols)
+ {
+ return SetOutputColumns(cols, partsupp_types_, partsupp_name_map_, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch()
+ {
+ size_t thread_index = thread_indexer_();
Review comment:
Yes, as long as the thread indexes within an exec plan are consistent. We should maybe have a dedicated discussion for this elsewhere, I think Michal also had some thoughts on something like this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r824953841
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string> &columns,
+ const std::vector<std::shared_ptr<DataType>> &types,
+ const std::unordered_map<std::string, int> &name_map,
+ std::vector<int> &gen_list)
+ {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if(columns.empty())
+ {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for(auto pair : name_map)
+ {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ }
+ else
+ {
+ for(const std::string &col : columns)
+ {
+ auto entry = name_map.find(col);
+ if(entry == name_map.end())
+ return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+ }
+
+ static TpchText g_text;
+
+ Status TpchText::InitIfNeeded(random::pcg32_fast &rng)
+ {
+ if(done_.load())
+ return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(!text_)
Review comment:
`text_` is uninitialized the first time this function is called.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r818061003
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,69 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
Review comment:
I don't think these headers (or the ones in vendored) are actually pulled in by default, at least they weren't pulled in when I built using the instructions: https://arrow.apache.org/docs/r/articles/developers/setup.html#step-2---configure-the-libarrow-build
I coppied the vendored/pcg/*.hpp over manually — but I suspect we'll need to add those to cmake somewhere to do that as part of the build now that these are used in compute
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r818092031
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,69 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
Review comment:
Good point. The `pcg` subdirectory needs to be added to https://github.com/apache/arrow/blob/master/cpp/src/arrow/vendored/CMakeLists.txt
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r817249386
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3738 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status Init();
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ void GenerateWord(size_t &offset, const char **words, size_t num_choices);
+ void GenerateNoun(size_t &offset);
+ void GenerateVerb(size_t &offset);
+ void GenerateAdjective(size_t &offset);
+ void GenerateAdverb(size_t &offset);
+ void GeneratePreposition(size_t &offset);
+ void GenerateAuxiliary(size_t &offset);
+ void GenerateTerminator(size_t &offset);
+
+ void GenerateNounPhrase(size_t &offset);
+ void GenerateVerbPhrase(size_t &offset);
+ void GeneratePrepositionalPhrase(size_t &offset);
+
+ void GenerateSentence(size_t &offset);
+
+ std::unique_ptr<Buffer> text_;
+ random::pcg32_fast rng_;
+ static constexpr size_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ int scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_generated_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
Review comment:
We should confirm that this is ok to copy/paste. https://www.apache.org/legal/resolved.html#stackoverflow says no, though I thought I remember mention of a date after which stack overflow content is ok.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082227200
Apologies, I'll not revert like that in the future. We can always revert the revert.
As for the problem of increasing transitive inclusions, I agree keeping compilation time down is a worthy goal. I guess this header file now will transitively include `pcg_random.h`. Maybe we can find some other way to avoid including `pcg_random` without having to make everything virtual.
I think either way we don't have to make the factory return a unique_ptr.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r818750831
##########
File path: r/src/compute-exec.cpp
##########
@@ -277,4 +287,135 @@ std::shared_ptr<compute::ExecNode> ExecNode_ReadFromRecordBatchReader(
return MakeExecNodeOrStop("source", plan.get(), {}, options);
}
+// [[arrow::export]]
+std::shared_ptr<arrow::RecordBatchReader> Tpch_Dbgen(
+ const std::shared_ptr<compute::ExecPlan>& plan,
+ int scale_factor,
+ std::string table_name
+ ) {
+
+ auto gen = ValueOrStop(arrow::compute::TpchGen::Make(plan.get(), scale_factor));
+
+ compute::ExecNode *table;
+ if (table_name == "part") {
+ table = ValueOrStop(gen.Part());
+ } else if (table_name == "supplier") {
+ table = ValueOrStop(gen.Supplier());
+ } else if (table_name == "partsupp") {
+ table = ValueOrStop(gen.PartSupp());
+ } else if (table_name == "customer") {
+ table = ValueOrStop(gen.Customer());
+ } else if (table_name == "nation") {
+ table = ValueOrStop(gen.Nation());
+ } else if (table_name == "lineitem") {
+ table = ValueOrStop(gen.Lineitem());
+ } else if (table_name == "region") {
+ table = ValueOrStop(gen.Region());
+ } else if (table_name == "orders") {
+ table = ValueOrStop(gen.Orders());
+ } else {
+ cpp11::stop("That's not a valid table name");
+ }
+
+ arrow::AsyncGenerator<arrow::util::optional<compute::ExecBatch>> sink_gen;
+
+ MakeExecNodeOrStop("sink", plan.get(), {table},
+ compute::SinkNodeOptions{&sink_gen});
+
+ StopIfNotOk(plan->Validate());
+ StopIfNotOk(plan->StartProducing());
+
+ // If the generator is destroyed before being completely drained, inform plan
+ std::shared_ptr<void> stop_producing{nullptr, [plan](...) {
+ bool not_finished_yet =
+ plan->finished().TryAddCallback([&plan] {
+ return [plan](const arrow::Status&) {};
+ });
+
+ if (not_finished_yet) {
+ plan->StopProducing();
+ }
+ }};
+
+ return compute::MakeGeneratorReader(
+ table->output_schema(),
+ [stop_producing, plan, sink_gen] { return sink_gen(); }, gc_memory_pool());
+}
+
+// [[arrow::export]]
+void Tpch_Dbgen_Write(
+ const std::shared_ptr<compute::ExecPlan>& plan,
+ int scale_factor,
+ std::string table_name,
+ const std::shared_ptr<fs::FileSystem>& filesystem, std::string base_dir,
+ arrow::dataset::ExistingDataBehavior existing_data_behavior, int max_partitions
+) {
+ auto gen = ValueOrStop(arrow::compute::TpchGen::Make(plan.get(), scale_factor));
+
+ compute::ExecNode *table;
+ if (table_name == "part") {
+ table = ValueOrStop(gen.Part());
+ } else if (table_name == "supplier") {
+ table = ValueOrStop(gen.Supplier());
+ } else if (table_name == "partsupp") {
+ table = ValueOrStop(gen.PartSupp());
+ } else if (table_name == "customer") {
+ table = ValueOrStop(gen.Customer());
+ } else if (table_name == "nation") {
+ table = ValueOrStop(gen.Nation());
+ } else if (table_name == "lineitem") {
+ table = ValueOrStop(gen.Lineitem());
+ } else if (table_name == "region") {
+ table = ValueOrStop(gen.Region());
+ } else if (table_name == "orders") {
+ table = ValueOrStop(gen.Orders());
+ } else {
+ cpp11::stop("That's not a valid table name");
+ }
+
+ // TODO: unhardcode this once it's working
+ auto base_path = base_dir + "/parquet_dataset";
+ filesystem->CreateDir(base_path);
+
+ auto format = std::make_shared<ds::ParquetFileFormat>();
+
+ ds::FileSystemDatasetWriteOptions write_options;
+ write_options.file_write_options = format->DefaultWriteOptions();
+ write_options.existing_data_behavior = ds::ExistingDataBehavior::kDeleteMatchingPartitions;
+ write_options.filesystem = filesystem;
+ write_options.base_dir = base_path;
+ write_options.partitioning = arrow::dataset::Partitioning::Default();
+ write_options.basename_template = "part{i}.parquet";
+ write_options.max_partitions = 1024;
+
+ // TODO: this had a checked_cast in front of it in the code I adapted it from
+ // but I ran into namespace issues when doing it so I took it out to see if it
+ // worked, but maybe that's what's causing the sefault?
+ const ds::WriteNodeOptions options =
+ ds::WriteNodeOptions{write_options, table->output_schema()};
Review comment:
This might be exactly what's causing the seg fault, the code I saw that did this had `checked_cast` in front of this
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r823395915
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3758 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ random::pcg32_fast rng_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ int scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_generated_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
Review comment:
Bisection was an interesting idea, I profiled it and it turned out to be slower than the linear chain of if statements. I'm not quite sure why that is, but I'd guess it has something to do with being harder for branch predictor.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r837459627
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,123 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <benchmark/benchmark.h>
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+namespace internal {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<double>(scale_factor));
+
+ ExecNode* lineitem =
+ *gen.Lineitem({"L_QUANTITY", "L_EXTENDEDPRICE", "L_TAX", "L_DISCOUNT", "L_SHIPDATE",
+ "L_RETURNFLAG", "L_LINESTATUS"});
+
+ auto sept_2_1998 = std::make_shared<Date32Scalar>(
+ 10471); // September 2, 1998 is 10471 days after January 1, 1970
+ Expression filter =
+ less_equal(field_ref("L_SHIPDATE"), literal(std::move(sept_2_1998)));
+ FilterNodeOptions filter_opts(filter);
+
+ Expression l_returnflag = field_ref("L_RETURNFLAG");
+ Expression l_linestatus = field_ref("L_LINESTATUS");
+ Expression quantity = field_ref("L_QUANTITY");
+ Expression base_price = field_ref("L_EXTENDEDPRICE");
+
+ std::shared_ptr<Decimal128Scalar> decimal_1 =
+ std::make_shared<Decimal128Scalar>(Decimal128{0, 100}, decimal(12, 2));
+ Expression discount_multiplier =
+ call("subtract", {literal(decimal_1), field_ref("L_DISCOUNT")});
+ Expression tax_multiplier = call("add", {literal(decimal_1), field_ref("L_TAX")});
+ Expression disc_price =
+ call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier});
+ Expression charge =
+ call("multiply",
+ {call("cast",
+ {call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier})},
+ compute::CastOptions::Unsafe(decimal(12, 2))),
+ tax_multiplier});
+ Expression discount = field_ref("L_DISCOUNT");
+
+ std::vector<Expression> projection_list = {l_returnflag, l_linestatus, quantity,
+ base_price, disc_price, charge,
+ quantity, base_price, discount};
+ std::vector<std::string> project_names = {
+ "l_returnflag", "l_linestatus", "sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price", "avg_disc"};
+ ProjectNodeOptions project_opts(std::move(projection_list), std::move(project_names));
+
+ ScalarAggregateOptions sum_opts = ScalarAggregateOptions::Defaults();
+ CountOptions count_opts(CountOptions::CountMode::ALL);
+ std::vector<arrow::compute::internal::Aggregate> aggs = {
+ {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts},
+ {"hash_sum", &sum_opts}, {"hash_mean", &sum_opts}, {"hash_mean", &sum_opts},
+ {"hash_mean", &sum_opts}, {"hash_count", &count_opts}};
+
+ std::vector<FieldRef> to_aggregate = {"sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price",
+ "avg_disc", "sum_qty"};
+
+ std::vector<std::string> names = {"sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price",
+ "avg_disc", "count_order"};
+
+ std::vector<FieldRef> keys = {"l_returnflag", "l_linestatus"};
+ AggregateNodeOptions agg_opts(aggs, to_aggregate, names, keys);
+
+ SortKey l_returnflag_key("l_returnflag");
+ SortKey l_linestatus_key("l_linestatus");
+ SortOptions sort_opts({l_returnflag_key, l_linestatus_key});
+ OrderBySinkNodeOptions order_by_opts(sort_opts, sink_gen);
+
+ Declaration filter_decl("filter", {Declaration::Input(lineitem)}, filter_opts);
+ Declaration project_decl("project", project_opts);
+ Declaration aggregate_decl("aggregate", agg_opts);
+ Declaration orderby_decl("order_by_sink", order_by_opts);
+
+ Declaration q1 =
+ Declaration::Sequence({filter_decl, project_decl, aggregate_decl, orderby_decl});
+ std::ignore = *q1.AddToPlan(plan.get());
+ return plan;
+}
+
+static void BM_Tpch_Q1(benchmark::State& st) {
+ for (auto _ : st) {
+ st.PauseTiming();
+ AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+ std::shared_ptr<ExecPlan> plan = Plan_Q1(&sink_gen, static_cast<int>(st.range(0)));
+ st.ResumeTiming();
+ auto fut = StartAndCollect(plan.get(), sink_gen);
+ auto res = *fut.MoveResult();
+ }
+}
Review comment:
Hmm, can't the metric simply be "queries per second"? Or does it not make sense for TPC-H?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082826737
Great! The s3fs failures on macOS are unrelated, so I'm gonna merge now. Thank you @save-buffer .
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1064176358
No, this is all original content. I followed the official spec, but no code was taken.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1067544853
Seems to be all green except for the JNI one which seems to be failing for everyone right now
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1077904240
the problem is that we actually don't know how many batches will be generated for LINEITEM because for each row in ORDERS there's a random number (between 1 and 7) of rows in LINEITEM.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r834554532
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row + num_offset;
+ int64_t num = 0;
+ int ibyte = static_cast<int>(num_offset);
+ for (; *num_str && ibyte < byte_width; ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ if (verify_padding) {
+ int num_chars = ibyte - num_offset;
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ VerifyStringAndNumber_Single(row, prefix, i, nums, byte_width, verify_padding);
+ }
+}
+
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ char tmp_str[256] = {};
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ std::memcpy(tmp_str, str + start, str_len);
+ VerifyStringAndNumber_Single(tmp_str, prefix, i, nums, sizeof(tmp_str),
+ verify_padding);
+ }
+}
+
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = off[i];
+ int32_t end = off[i + 1];
+ int32_t str_len = end - start;
+ ASSERT_LE(str_len, max_length);
+ ASSERT_GE(str_len, min_length);
+ for (int32_t i = start; i < end; i++) {
+ bool is_valid =
+ std::isdigit(str[i]) || std::isalpha(str[i]) || str[i] == ',' || str[i] == ' ';
+ ASSERT_TRUE(is_valid) << "Character " << str[i]
+ << " is not a digit, a letter, a comma, or a space";
+ }
+ }
+}
+
+void VerifyModuloBetween(const Datum& d, int32_t min, int32_t max, int32_t mod) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t m = n[i] % mod;
+ ASSERT_GE(m, min) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ ASSERT_LE(m, max) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ }
+}
+
+void VerifyAllBetween(const Datum& d, int32_t min, int32_t max) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ ASSERT_GE(n[i], min) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ ASSERT_LE(n[i], max) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ }
+}
+
+void VerifyNationKey(const Datum& d) { VerifyAllBetween(d, 0, 24); }
+
+void VerifyPhone(const Datum& d) {
+ int64_t length = d.length();
+ const char* phones = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ constexpr int kByteWidth = 15; // This is common for all PHONE columns
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = phones + i * kByteWidth;
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ }
+}
+
+void VerifyDecimalsBetween(const Datum& d, int64_t min, int64_t max) {
+ int64_t length = d.length();
+ const Decimal128* decs =
+ reinterpret_cast<const Decimal128*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int64_t val = static_cast<int64_t>(decs[i]);
+ ASSERT_LE(val, max);
+ ASSERT_GE(val, min);
+ }
+}
+
+void VerifyCorrectNumberOfWords_Varlen(const Datum& d, int num_words) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+
+ int32_t start = offsets[i];
+ int32_t end = offsets[i + 1];
+ int32_t str_len = end - start;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, str + start, str_len);
Review comment:
sorry about that, I had fixed a different occurrence of this pattern elsewhere. i guess i had three, i've fixed the other two.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1078665383
Looks like the failing checks are due to hash-join-node-test timing out and
```
unknown file: error: SEH exception with code 0xc00000fd thrown in the test body.
[ FAILED ] TestScan/TestCsvFileFormatScan.ScanBatchSize/0Threaded16b1024r, where GetParam() = Threaded16b1024r (218 ms)
```
Don't think it's this PR's problem
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1077319605
@pitrou did you want to take another look at this?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r834589559
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,631 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <regex>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+namespace internal {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+// Verifies that the data is valid Arrow and ensures it's not null.
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values) {
+ ASSERT_EQ(d.array()->buffers[0].get(), nullptr);
+ ASSERT_OK(d.make_array()->ValidateFull());
+ }
+}
+
+// Verifies that each element is seen exactly once, and that it's between min and max
+// inclusive
+void VerifyUniqueKey(std::unordered_set<int32_t>* seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen->insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const util::string_view& row, const char* prefix,
+ const int64_t i, const int32_t* nums,
+ bool verify_padding) {
+ size_t num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_TRUE(row.starts_with(prefix)) << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row.data() + num_offset;
+ int64_t num = 0;
+ size_t ibyte = num_offset;
+ // Parse the number out
+ for (; *num_str && ibyte < row.size(); ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ // If nums is not null, ensure it matches the parsed number
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ // TPC-H requires only ever requires padding up to 9 digits, so we ensure that
+ // the total length of the string was at least 9 (could be more for bigger numbers).
+ if (verify_padding) {
+ int64_t num_chars = static_cast<int64_t>(ibyte - num_offset);
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+// Verifies that each row is the string "prefix" followed by a number. If numbers is not
+// EMPTY, it also checks that the number following the prefix is equal to the
+// corresponding row in numbers. Some TPC-H data is padded to 9 zeros, which this function
+// can optionally verify as well. This string function verifies fixed width columns.
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ util::string_view view(row, byte_width);
+ VerifyStringAndNumber_Single(view, prefix, i, nums, verify_padding);
+ }
+}
+
+// Same as above but for variable length columns
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ util::string_view view(str + start, str_len);
+ VerifyStringAndNumber_Single(view, prefix, i, nums, verify_padding);
+ }
+}
+
+// Verifies that each row is a V-string, which is defined in the spec to be
+// a string of random length between min_length and max_length, that is composed
+// of alphanumeric characters, commas, or spaces.
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = off[i];
+ int32_t end = off[i + 1];
+ int32_t str_len = end - start;
+ ASSERT_LE(str_len, max_length);
+ ASSERT_GE(str_len, min_length);
+ for (int32_t i = start; i < end; i++) {
+ bool is_valid =
+ std::isdigit(str[i]) || std::isalpha(str[i]) || str[i] == ',' || str[i] == ' ';
+ ASSERT_TRUE(is_valid) << "Character " << str[i]
+ << " is not a digit, a letter, a comma, or a space";
+ }
+ }
+}
+
+// Verifies that each 32-bit element modulo "mod" is between min and max.
+void VerifyModuloBetween(const Datum& d, int32_t min, int32_t max, int32_t mod) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t m = n[i] % mod;
+ ASSERT_GE(m, min) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ ASSERT_LE(m, max) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ }
+}
+
+// Verifies that each 32-bit element is between min and max.
+void VerifyAllBetween(const Datum& d, int32_t min, int32_t max) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ ASSERT_GE(n[i], min) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ ASSERT_LE(n[i], max) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ }
+}
+
+void VerifyNationKey(const Datum& d) { VerifyAllBetween(d, 0, 24); }
+
+// Verifies that each row satisfies the phone number spec.
+void VerifyPhone(const Datum& d) {
+ int64_t length = d.length();
+ const char* phones = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ constexpr int kByteWidth = 15; // This is common for all PHONE columns
+ std::regex exp("\\d{2}-\\d{3}-\\d{3}-\\d{4}");
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = phones + i * kByteWidth;
+ ASSERT_TRUE(std::regex_match(row, row + kByteWidth, exp));
+ }
+}
+
+// Verifies that each decimal is between min and max
+void VerifyDecimalsBetween(const Datum& d, int64_t min, int64_t max) {
+ int64_t length = d.length();
+ const Decimal128* decs =
+ reinterpret_cast<const Decimal128*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int64_t val = static_cast<int64_t>(decs[i]);
+ ASSERT_LE(val, max);
+ ASSERT_GE(val, min);
+ }
+}
+
+// Verifies that each variable-length row is a series of words separated by
+// spaces. Number of words is determined by the number of spaces.
+void VerifyCorrectNumberOfWords_Varlen(const Datum& d, int num_words) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+
+ int32_t start = offsets[i];
+ int32_t end = offsets[i + 1];
+ int32_t str_len = end - start;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, str + start, str_len);
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = offsets[i]; j < offsets[i + 1]; j++) {
+ bool is_space = str[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(str[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << tmp_str;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << tmp_str;
+ }
+}
+
+// Same as above but for fixed width columns.
+void VerifyCorrectNumberOfWords_FixedWidth(const Datum& d, int num_words,
+ int byte_width) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+ const char* row = str + i * byte_width;
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = 0; j < byte_width && row[j]; j++) {
+ bool is_space = row[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(row[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << row;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << row;
+ }
+}
+
+// Verifies that each row of the single-byte-wide column is one of the possibilities.
+void VerifyOneOf(const Datum& d, const std::unordered_set<char>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++)
+ ASSERT_TRUE(possibilities.find(col[i]) != possibilities.end());
+}
+
+// Verifies that each fixed-width row is one of the possibilities
+void VerifyOneOf(const Datum& d, int32_t byte_width,
+ const std::unordered_set<util::string_view>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = col + i * byte_width;
+ int32_t row_len = 0;
+ while (row[row_len] && row_len < byte_width) row_len++;
+ util::string_view view(row, row_len);
+ ASSERT_TRUE(possibilities.find(view) != possibilities.end())
+ << view << " is not a valid string.";
+ }
+}
+
+// Counts the number of instances of each integer
+void CountInstances(std::unordered_map<int32_t, int32_t>& counts, const Datum& d) {
+ int64_t length = d.length();
+ const int32_t* nums = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) counts[nums[i]]++;
+}
+
+// For the S_COMMENT column, some of the columns must be modified to contain
+// "Customer...Complaints" or "Customer...Recommends". This function counts the number of
+// good and bad comments.
+void CountModifiedComments(const Datum& d, int& good_count, int& bad_count) {
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ // Length of S_COMMENT is at most 100
+ char tmp_string[101];
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + offsets[i];
+ int32_t row_length = offsets[i + 1] - offsets[i];
+ std::memset(tmp_string, 0, sizeof(tmp_string));
+ std::memcpy(tmp_string, row, row_length);
+ char* customer = std::strstr(tmp_string, "Customer");
+ char* recommends = std::strstr(tmp_string, "Recommends");
+ char* complaints = std::strstr(tmp_string, "Complaints");
+ if (customer) {
+ ASSERT_TRUE((recommends != nullptr) ^ (complaints != nullptr));
+ if (recommends) good_count++;
+ if (complaints) bad_count++;
+ }
+ }
+}
+
+TEST(TpchNode, ScaleFactor) {
+ ExecContext ctx(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(&ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), 0.25f);
+ ExecNode* table = *gen.Supplier();
+ AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+ Declaration sink("sink", {Declaration::Input(table)}, SinkNodeOptions{&sink_gen});
+ std::ignore = *sink.AddToPlan(plan.get());
+ auto fut = StartAndCollect(plan.get(), sink_gen);
+ auto res = *fut.MoveResult();
Review comment:
i've factored this setup code into a function, so i think i can just get away with `return fut.MoveResult()` because the function returns a `Result`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r834030251
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,631 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <regex>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+namespace internal {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+// Verifies that the data is valid Arrow and ensures it's not null.
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values) {
+ ASSERT_EQ(d.array()->buffers[0].get(), nullptr);
+ ASSERT_OK(d.make_array()->ValidateFull());
+ }
+}
+
+// Verifies that each element is seen exactly once, and that it's between min and max
+// inclusive
+void VerifyUniqueKey(std::unordered_set<int32_t>* seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen->insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const util::string_view& row, const char* prefix,
+ const int64_t i, const int32_t* nums,
+ bool verify_padding) {
+ size_t num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_TRUE(row.starts_with(prefix)) << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row.data() + num_offset;
+ int64_t num = 0;
+ size_t ibyte = num_offset;
+ // Parse the number out
+ for (; *num_str && ibyte < row.size(); ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ // If nums is not null, ensure it matches the parsed number
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ // TPC-H requires only ever requires padding up to 9 digits, so we ensure that
+ // the total length of the string was at least 9 (could be more for bigger numbers).
+ if (verify_padding) {
+ int64_t num_chars = static_cast<int64_t>(ibyte - num_offset);
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+// Verifies that each row is the string "prefix" followed by a number. If numbers is not
+// EMPTY, it also checks that the number following the prefix is equal to the
+// corresponding row in numbers. Some TPC-H data is padded to 9 zeros, which this function
+// can optionally verify as well. This string function verifies fixed width columns.
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ util::string_view view(row, byte_width);
+ VerifyStringAndNumber_Single(view, prefix, i, nums, verify_padding);
+ }
+}
+
+// Same as above but for variable length columns
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ util::string_view view(str + start, str_len);
+ VerifyStringAndNumber_Single(view, prefix, i, nums, verify_padding);
+ }
+}
+
+// Verifies that each row is a V-string, which is defined in the spec to be
+// a string of random length between min_length and max_length, that is composed
+// of alphanumeric characters, commas, or spaces.
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = off[i];
+ int32_t end = off[i + 1];
+ int32_t str_len = end - start;
+ ASSERT_LE(str_len, max_length);
+ ASSERT_GE(str_len, min_length);
+ for (int32_t i = start; i < end; i++) {
+ bool is_valid =
+ std::isdigit(str[i]) || std::isalpha(str[i]) || str[i] == ',' || str[i] == ' ';
+ ASSERT_TRUE(is_valid) << "Character " << str[i]
+ << " is not a digit, a letter, a comma, or a space";
+ }
+ }
+}
+
+// Verifies that each 32-bit element modulo "mod" is between min and max.
+void VerifyModuloBetween(const Datum& d, int32_t min, int32_t max, int32_t mod) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t m = n[i] % mod;
+ ASSERT_GE(m, min) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ ASSERT_LE(m, max) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ }
+}
+
+// Verifies that each 32-bit element is between min and max.
+void VerifyAllBetween(const Datum& d, int32_t min, int32_t max) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ ASSERT_GE(n[i], min) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ ASSERT_LE(n[i], max) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ }
+}
+
+void VerifyNationKey(const Datum& d) { VerifyAllBetween(d, 0, 24); }
+
+// Verifies that each row satisfies the phone number spec.
+void VerifyPhone(const Datum& d) {
+ int64_t length = d.length();
+ const char* phones = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ constexpr int kByteWidth = 15; // This is common for all PHONE columns
+ std::regex exp("\\d{2}-\\d{3}-\\d{3}-\\d{4}");
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = phones + i * kByteWidth;
+ ASSERT_TRUE(std::regex_match(row, row + kByteWidth, exp));
+ }
+}
+
+// Verifies that each decimal is between min and max
+void VerifyDecimalsBetween(const Datum& d, int64_t min, int64_t max) {
+ int64_t length = d.length();
+ const Decimal128* decs =
+ reinterpret_cast<const Decimal128*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int64_t val = static_cast<int64_t>(decs[i]);
+ ASSERT_LE(val, max);
+ ASSERT_GE(val, min);
+ }
+}
+
+// Verifies that each variable-length row is a series of words separated by
+// spaces. Number of words is determined by the number of spaces.
+void VerifyCorrectNumberOfWords_Varlen(const Datum& d, int num_words) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+
+ int32_t start = offsets[i];
+ int32_t end = offsets[i + 1];
+ int32_t str_len = end - start;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, str + start, str_len);
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = offsets[i]; j < offsets[i + 1]; j++) {
+ bool is_space = str[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(str[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << tmp_str;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << tmp_str;
+ }
+}
+
+// Same as above but for fixed width columns.
+void VerifyCorrectNumberOfWords_FixedWidth(const Datum& d, int num_words,
+ int byte_width) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+ const char* row = str + i * byte_width;
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = 0; j < byte_width && row[j]; j++) {
+ bool is_space = row[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(row[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << row;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << row;
+ }
+}
+
+// Verifies that each row of the single-byte-wide column is one of the possibilities.
+void VerifyOneOf(const Datum& d, const std::unordered_set<char>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++)
+ ASSERT_TRUE(possibilities.find(col[i]) != possibilities.end());
+}
+
+// Verifies that each fixed-width row is one of the possibilities
+void VerifyOneOf(const Datum& d, int32_t byte_width,
+ const std::unordered_set<util::string_view>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = col + i * byte_width;
+ int32_t row_len = 0;
+ while (row[row_len] && row_len < byte_width) row_len++;
+ util::string_view view(row, row_len);
+ ASSERT_TRUE(possibilities.find(view) != possibilities.end())
+ << view << " is not a valid string.";
+ }
+}
+
+// Counts the number of instances of each integer
+void CountInstances(std::unordered_map<int32_t, int32_t>& counts, const Datum& d) {
+ int64_t length = d.length();
+ const int32_t* nums = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) counts[nums[i]]++;
+}
+
+// For the S_COMMENT column, some of the columns must be modified to contain
+// "Customer...Complaints" or "Customer...Recommends". This function counts the number of
+// good and bad comments.
+void CountModifiedComments(const Datum& d, int& good_count, int& bad_count) {
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ // Length of S_COMMENT is at most 100
+ char tmp_string[101];
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + offsets[i];
+ int32_t row_length = offsets[i + 1] - offsets[i];
+ std::memset(tmp_string, 0, sizeof(tmp_string));
+ std::memcpy(tmp_string, row, row_length);
+ char* customer = std::strstr(tmp_string, "Customer");
+ char* recommends = std::strstr(tmp_string, "Recommends");
+ char* complaints = std::strstr(tmp_string, "Complaints");
+ if (customer) {
+ ASSERT_TRUE((recommends != nullptr) ^ (complaints != nullptr));
+ if (recommends) good_count++;
+ if (complaints) bad_count++;
+ }
+ }
+}
+
+TEST(TpchNode, ScaleFactor) {
+ ExecContext ctx(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(&ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), 0.25f);
+ ExecNode* table = *gen.Supplier();
+ AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+ Declaration sink("sink", {Declaration::Input(table)}, SinkNodeOptions{&sink_gen});
+ std::ignore = *sink.AddToPlan(plan.get());
+ auto fut = StartAndCollect(plan.get(), sink_gen);
+ auto res = *fut.MoveResult();
Review comment:
In unit tests, we shouldn't just crash the process if there is an unexpected error, but check for it.
```suggestion
ASSERT_FINISHES_OK_AND_ASSIGN(auto res, fut);
```
(note you'll need to include `arrow/testing/future_util.h`)
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,631 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <regex>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+namespace internal {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+// Verifies that the data is valid Arrow and ensures it's not null.
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values) {
+ ASSERT_EQ(d.array()->buffers[0].get(), nullptr);
+ ASSERT_OK(d.make_array()->ValidateFull());
+ }
+}
+
+// Verifies that each element is seen exactly once, and that it's between min and max
+// inclusive
+void VerifyUniqueKey(std::unordered_set<int32_t>* seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen->insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const util::string_view& row, const char* prefix,
+ const int64_t i, const int32_t* nums,
+ bool verify_padding) {
+ size_t num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_TRUE(row.starts_with(prefix)) << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row.data() + num_offset;
+ int64_t num = 0;
+ size_t ibyte = num_offset;
+ // Parse the number out
+ for (; *num_str && ibyte < row.size(); ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ // If nums is not null, ensure it matches the parsed number
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ // TPC-H requires only ever requires padding up to 9 digits, so we ensure that
+ // the total length of the string was at least 9 (could be more for bigger numbers).
+ if (verify_padding) {
+ int64_t num_chars = static_cast<int64_t>(ibyte - num_offset);
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+// Verifies that each row is the string "prefix" followed by a number. If numbers is not
+// EMPTY, it also checks that the number following the prefix is equal to the
+// corresponding row in numbers. Some TPC-H data is padded to 9 zeros, which this function
+// can optionally verify as well. This string function verifies fixed width columns.
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ util::string_view view(row, byte_width);
+ VerifyStringAndNumber_Single(view, prefix, i, nums, verify_padding);
+ }
+}
+
+// Same as above but for variable length columns
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ util::string_view view(str + start, str_len);
+ VerifyStringAndNumber_Single(view, prefix, i, nums, verify_padding);
+ }
+}
+
+// Verifies that each row is a V-string, which is defined in the spec to be
+// a string of random length between min_length and max_length, that is composed
+// of alphanumeric characters, commas, or spaces.
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = off[i];
+ int32_t end = off[i + 1];
+ int32_t str_len = end - start;
+ ASSERT_LE(str_len, max_length);
+ ASSERT_GE(str_len, min_length);
+ for (int32_t i = start; i < end; i++) {
+ bool is_valid =
+ std::isdigit(str[i]) || std::isalpha(str[i]) || str[i] == ',' || str[i] == ' ';
+ ASSERT_TRUE(is_valid) << "Character " << str[i]
+ << " is not a digit, a letter, a comma, or a space";
+ }
+ }
+}
+
+// Verifies that each 32-bit element modulo "mod" is between min and max.
+void VerifyModuloBetween(const Datum& d, int32_t min, int32_t max, int32_t mod) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t m = n[i] % mod;
+ ASSERT_GE(m, min) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ ASSERT_LE(m, max) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ }
+}
+
+// Verifies that each 32-bit element is between min and max.
+void VerifyAllBetween(const Datum& d, int32_t min, int32_t max) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ ASSERT_GE(n[i], min) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ ASSERT_LE(n[i], max) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ }
+}
+
+void VerifyNationKey(const Datum& d) { VerifyAllBetween(d, 0, 24); }
+
+// Verifies that each row satisfies the phone number spec.
+void VerifyPhone(const Datum& d) {
+ int64_t length = d.length();
+ const char* phones = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ constexpr int kByteWidth = 15; // This is common for all PHONE columns
+ std::regex exp("\\d{2}-\\d{3}-\\d{3}-\\d{4}");
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = phones + i * kByteWidth;
+ ASSERT_TRUE(std::regex_match(row, row + kByteWidth, exp));
+ }
+}
+
+// Verifies that each decimal is between min and max
+void VerifyDecimalsBetween(const Datum& d, int64_t min, int64_t max) {
+ int64_t length = d.length();
+ const Decimal128* decs =
+ reinterpret_cast<const Decimal128*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int64_t val = static_cast<int64_t>(decs[i]);
+ ASSERT_LE(val, max);
+ ASSERT_GE(val, min);
+ }
+}
+
+// Verifies that each variable-length row is a series of words separated by
+// spaces. Number of words is determined by the number of spaces.
+void VerifyCorrectNumberOfWords_Varlen(const Datum& d, int num_words) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+
+ int32_t start = offsets[i];
+ int32_t end = offsets[i + 1];
+ int32_t str_len = end - start;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, str + start, str_len);
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = offsets[i]; j < offsets[i + 1]; j++) {
+ bool is_space = str[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(str[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << tmp_str;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << tmp_str;
+ }
+}
+
+// Same as above but for fixed width columns.
+void VerifyCorrectNumberOfWords_FixedWidth(const Datum& d, int num_words,
+ int byte_width) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+ const char* row = str + i * byte_width;
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = 0; j < byte_width && row[j]; j++) {
+ bool is_space = row[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(row[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << row;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << row;
+ }
+}
+
+// Verifies that each row of the single-byte-wide column is one of the possibilities.
+void VerifyOneOf(const Datum& d, const std::unordered_set<char>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++)
+ ASSERT_TRUE(possibilities.find(col[i]) != possibilities.end());
+}
+
+// Verifies that each fixed-width row is one of the possibilities
+void VerifyOneOf(const Datum& d, int32_t byte_width,
+ const std::unordered_set<util::string_view>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = col + i * byte_width;
+ int32_t row_len = 0;
+ while (row[row_len] && row_len < byte_width) row_len++;
+ util::string_view view(row, row_len);
+ ASSERT_TRUE(possibilities.find(view) != possibilities.end())
+ << view << " is not a valid string.";
+ }
+}
+
+// Counts the number of instances of each integer
+void CountInstances(std::unordered_map<int32_t, int32_t>& counts, const Datum& d) {
+ int64_t length = d.length();
+ const int32_t* nums = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) counts[nums[i]]++;
+}
+
+// For the S_COMMENT column, some of the columns must be modified to contain
+// "Customer...Complaints" or "Customer...Recommends". This function counts the number of
+// good and bad comments.
+void CountModifiedComments(const Datum& d, int& good_count, int& bad_count) {
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ // Length of S_COMMENT is at most 100
+ char tmp_string[101];
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + offsets[i];
+ int32_t row_length = offsets[i + 1] - offsets[i];
+ std::memset(tmp_string, 0, sizeof(tmp_string));
+ std::memcpy(tmp_string, row, row_length);
+ char* customer = std::strstr(tmp_string, "Customer");
+ char* recommends = std::strstr(tmp_string, "Recommends");
+ char* complaints = std::strstr(tmp_string, "Complaints");
+ if (customer) {
+ ASSERT_TRUE((recommends != nullptr) ^ (complaints != nullptr));
+ if (recommends) good_count++;
+ if (complaints) bad_count++;
+ }
+ }
+}
+
+TEST(TpchNode, ScaleFactor) {
+ ExecContext ctx(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(&ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), 0.25f);
+ ExecNode* table = *gen.Supplier();
+ AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+ Declaration sink("sink", {Declaration::Input(table)}, SinkNodeOptions{&sink_gen});
+ std::ignore = *sink.AddToPlan(plan.get());
+ auto fut = StartAndCollect(plan.get(), sink_gen);
+ auto res = *fut.MoveResult();
+
+ int64_t kExpectedRows = 2500;
+ int64_t num_rows = 0;
+ for (auto& batch : res) num_rows += batch.length;
+ ASSERT_EQ(num_rows, kExpectedRows);
+}
+
+TEST(TpchNode, Supplier) {
+ ExecContext ctx(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(&ctx);
+ TpchGen gen = *TpchGen::Make(plan.get());
+ ExecNode* table = *gen.Supplier();
+ AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+ Declaration sink("sink", {Declaration::Input(table)}, SinkNodeOptions{&sink_gen});
+ std::ignore = *sink.AddToPlan(plan.get());
+ auto fut = StartAndCollect(plan.get(), sink_gen);
+ auto res = *fut.MoveResult();
Review comment:
Same here as well of course, and in other instances below.
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,631 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <regex>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+namespace internal {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+// Verifies that the data is valid Arrow and ensures it's not null.
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values) {
+ ASSERT_EQ(d.array()->buffers[0].get(), nullptr);
+ ASSERT_OK(d.make_array()->ValidateFull());
+ }
+}
+
+// Verifies that each element is seen exactly once, and that it's between min and max
+// inclusive
+void VerifyUniqueKey(std::unordered_set<int32_t>* seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen->insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const util::string_view& row, const char* prefix,
+ const int64_t i, const int32_t* nums,
+ bool verify_padding) {
+ size_t num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_TRUE(row.starts_with(prefix)) << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row.data() + num_offset;
+ int64_t num = 0;
+ size_t ibyte = num_offset;
+ // Parse the number out
+ for (; *num_str && ibyte < row.size(); ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ // If nums is not null, ensure it matches the parsed number
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ // TPC-H requires only ever requires padding up to 9 digits, so we ensure that
+ // the total length of the string was at least 9 (could be more for bigger numbers).
+ if (verify_padding) {
+ int64_t num_chars = static_cast<int64_t>(ibyte - num_offset);
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+// Verifies that each row is the string "prefix" followed by a number. If numbers is not
+// EMPTY, it also checks that the number following the prefix is equal to the
+// corresponding row in numbers. Some TPC-H data is padded to 9 zeros, which this function
+// can optionally verify as well. This string function verifies fixed width columns.
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ util::string_view view(row, byte_width);
+ VerifyStringAndNumber_Single(view, prefix, i, nums, verify_padding);
+ }
+}
+
+// Same as above but for variable length columns
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ util::string_view view(str + start, str_len);
+ VerifyStringAndNumber_Single(view, prefix, i, nums, verify_padding);
+ }
+}
+
+// Verifies that each row is a V-string, which is defined in the spec to be
+// a string of random length between min_length and max_length, that is composed
+// of alphanumeric characters, commas, or spaces.
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = off[i];
+ int32_t end = off[i + 1];
+ int32_t str_len = end - start;
+ ASSERT_LE(str_len, max_length);
+ ASSERT_GE(str_len, min_length);
+ for (int32_t i = start; i < end; i++) {
+ bool is_valid =
+ std::isdigit(str[i]) || std::isalpha(str[i]) || str[i] == ',' || str[i] == ' ';
+ ASSERT_TRUE(is_valid) << "Character " << str[i]
+ << " is not a digit, a letter, a comma, or a space";
+ }
+ }
+}
+
+// Verifies that each 32-bit element modulo "mod" is between min and max.
+void VerifyModuloBetween(const Datum& d, int32_t min, int32_t max, int32_t mod) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t m = n[i] % mod;
+ ASSERT_GE(m, min) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ ASSERT_LE(m, max) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ }
+}
+
+// Verifies that each 32-bit element is between min and max.
+void VerifyAllBetween(const Datum& d, int32_t min, int32_t max) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ ASSERT_GE(n[i], min) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ ASSERT_LE(n[i], max) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ }
+}
+
+void VerifyNationKey(const Datum& d) { VerifyAllBetween(d, 0, 24); }
+
+// Verifies that each row satisfies the phone number spec.
+void VerifyPhone(const Datum& d) {
+ int64_t length = d.length();
+ const char* phones = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ constexpr int kByteWidth = 15; // This is common for all PHONE columns
+ std::regex exp("\\d{2}-\\d{3}-\\d{3}-\\d{4}");
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = phones + i * kByteWidth;
+ ASSERT_TRUE(std::regex_match(row, row + kByteWidth, exp));
+ }
+}
+
+// Verifies that each decimal is between min and max
+void VerifyDecimalsBetween(const Datum& d, int64_t min, int64_t max) {
+ int64_t length = d.length();
+ const Decimal128* decs =
+ reinterpret_cast<const Decimal128*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int64_t val = static_cast<int64_t>(decs[i]);
+ ASSERT_LE(val, max);
+ ASSERT_GE(val, min);
+ }
+}
+
+// Verifies that each variable-length row is a series of words separated by
+// spaces. Number of words is determined by the number of spaces.
+void VerifyCorrectNumberOfWords_Varlen(const Datum& d, int num_words) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+
+ int32_t start = offsets[i];
+ int32_t end = offsets[i + 1];
+ int32_t str_len = end - start;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, str + start, str_len);
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = offsets[i]; j < offsets[i + 1]; j++) {
+ bool is_space = str[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(str[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << tmp_str;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << tmp_str;
+ }
+}
+
+// Same as above but for fixed width columns.
+void VerifyCorrectNumberOfWords_FixedWidth(const Datum& d, int num_words,
+ int byte_width) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+ const char* row = str + i * byte_width;
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = 0; j < byte_width && row[j]; j++) {
+ bool is_space = row[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(row[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << row;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << row;
+ }
+}
+
+// Verifies that each row of the single-byte-wide column is one of the possibilities.
+void VerifyOneOf(const Datum& d, const std::unordered_set<char>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++)
+ ASSERT_TRUE(possibilities.find(col[i]) != possibilities.end());
+}
+
+// Verifies that each fixed-width row is one of the possibilities
+void VerifyOneOf(const Datum& d, int32_t byte_width,
+ const std::unordered_set<util::string_view>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = col + i * byte_width;
+ int32_t row_len = 0;
+ while (row[row_len] && row_len < byte_width) row_len++;
+ util::string_view view(row, row_len);
+ ASSERT_TRUE(possibilities.find(view) != possibilities.end())
+ << view << " is not a valid string.";
+ }
+}
+
+// Counts the number of instances of each integer
+void CountInstances(std::unordered_map<int32_t, int32_t>& counts, const Datum& d) {
Review comment:
Just like above, `counts` should be passed by pointer since it's an in-out parameter (also same comment for potential other occurrences of this pattern below).
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,631 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <regex>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+namespace internal {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+// Verifies that the data is valid Arrow and ensures it's not null.
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values) {
+ ASSERT_EQ(d.array()->buffers[0].get(), nullptr);
+ ASSERT_OK(d.make_array()->ValidateFull());
+ }
+}
+
+// Verifies that each element is seen exactly once, and that it's between min and max
+// inclusive
+void VerifyUniqueKey(std::unordered_set<int32_t>* seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen->insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const util::string_view& row, const char* prefix,
+ const int64_t i, const int32_t* nums,
+ bool verify_padding) {
+ size_t num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_TRUE(row.starts_with(prefix)) << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row.data() + num_offset;
+ int64_t num = 0;
+ size_t ibyte = num_offset;
+ // Parse the number out
+ for (; *num_str && ibyte < row.size(); ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ // If nums is not null, ensure it matches the parsed number
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ // TPC-H requires only ever requires padding up to 9 digits, so we ensure that
+ // the total length of the string was at least 9 (could be more for bigger numbers).
+ if (verify_padding) {
+ int64_t num_chars = static_cast<int64_t>(ibyte - num_offset);
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+// Verifies that each row is the string "prefix" followed by a number. If numbers is not
+// EMPTY, it also checks that the number following the prefix is equal to the
+// corresponding row in numbers. Some TPC-H data is padded to 9 zeros, which this function
+// can optionally verify as well. This string function verifies fixed width columns.
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ util::string_view view(row, byte_width);
+ VerifyStringAndNumber_Single(view, prefix, i, nums, verify_padding);
+ }
+}
+
+// Same as above but for variable length columns
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ util::string_view view(str + start, str_len);
+ VerifyStringAndNumber_Single(view, prefix, i, nums, verify_padding);
+ }
+}
+
+// Verifies that each row is a V-string, which is defined in the spec to be
+// a string of random length between min_length and max_length, that is composed
+// of alphanumeric characters, commas, or spaces.
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = off[i];
+ int32_t end = off[i + 1];
+ int32_t str_len = end - start;
+ ASSERT_LE(str_len, max_length);
+ ASSERT_GE(str_len, min_length);
+ for (int32_t i = start; i < end; i++) {
+ bool is_valid =
+ std::isdigit(str[i]) || std::isalpha(str[i]) || str[i] == ',' || str[i] == ' ';
+ ASSERT_TRUE(is_valid) << "Character " << str[i]
+ << " is not a digit, a letter, a comma, or a space";
+ }
+ }
+}
+
+// Verifies that each 32-bit element modulo "mod" is between min and max.
+void VerifyModuloBetween(const Datum& d, int32_t min, int32_t max, int32_t mod) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t m = n[i] % mod;
+ ASSERT_GE(m, min) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ ASSERT_LE(m, max) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ }
+}
+
+// Verifies that each 32-bit element is between min and max.
+void VerifyAllBetween(const Datum& d, int32_t min, int32_t max) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ ASSERT_GE(n[i], min) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ ASSERT_LE(n[i], max) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ }
+}
+
+void VerifyNationKey(const Datum& d) { VerifyAllBetween(d, 0, 24); }
+
+// Verifies that each row satisfies the phone number spec.
+void VerifyPhone(const Datum& d) {
+ int64_t length = d.length();
+ const char* phones = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ constexpr int kByteWidth = 15; // This is common for all PHONE columns
+ std::regex exp("\\d{2}-\\d{3}-\\d{3}-\\d{4}");
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = phones + i * kByteWidth;
+ ASSERT_TRUE(std::regex_match(row, row + kByteWidth, exp));
+ }
+}
+
+// Verifies that each decimal is between min and max
+void VerifyDecimalsBetween(const Datum& d, int64_t min, int64_t max) {
+ int64_t length = d.length();
+ const Decimal128* decs =
+ reinterpret_cast<const Decimal128*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int64_t val = static_cast<int64_t>(decs[i]);
+ ASSERT_LE(val, max);
+ ASSERT_GE(val, min);
+ }
+}
+
+// Verifies that each variable-length row is a series of words separated by
+// spaces. Number of words is determined by the number of spaces.
+void VerifyCorrectNumberOfWords_Varlen(const Datum& d, int num_words) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+
+ int32_t start = offsets[i];
+ int32_t end = offsets[i + 1];
+ int32_t str_len = end - start;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, str + start, str_len);
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = offsets[i]; j < offsets[i + 1]; j++) {
+ bool is_space = str[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(str[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << tmp_str;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << tmp_str;
+ }
+}
+
+// Same as above but for fixed width columns.
+void VerifyCorrectNumberOfWords_FixedWidth(const Datum& d, int num_words,
+ int byte_width) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+ const char* row = str + i * byte_width;
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = 0; j < byte_width && row[j]; j++) {
+ bool is_space = row[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(row[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << row;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << row;
+ }
+}
+
+// Verifies that each row of the single-byte-wide column is one of the possibilities.
+void VerifyOneOf(const Datum& d, const std::unordered_set<char>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++)
+ ASSERT_TRUE(possibilities.find(col[i]) != possibilities.end());
+}
+
+// Verifies that each fixed-width row is one of the possibilities
+void VerifyOneOf(const Datum& d, int32_t byte_width,
+ const std::unordered_set<util::string_view>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = col + i * byte_width;
+ int32_t row_len = 0;
+ while (row[row_len] && row_len < byte_width) row_len++;
+ util::string_view view(row, row_len);
+ ASSERT_TRUE(possibilities.find(view) != possibilities.end())
+ << view << " is not a valid string.";
+ }
+}
+
+// Counts the number of instances of each integer
+void CountInstances(std::unordered_map<int32_t, int32_t>& counts, const Datum& d) {
+ int64_t length = d.length();
+ const int32_t* nums = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) counts[nums[i]]++;
+}
+
+// For the S_COMMENT column, some of the columns must be modified to contain
+// "Customer...Complaints" or "Customer...Recommends". This function counts the number of
+// good and bad comments.
+void CountModifiedComments(const Datum& d, int& good_count, int& bad_count) {
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ // Length of S_COMMENT is at most 100
+ char tmp_string[101];
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + offsets[i];
+ int32_t row_length = offsets[i + 1] - offsets[i];
+ std::memset(tmp_string, 0, sizeof(tmp_string));
+ std::memcpy(tmp_string, row, row_length);
+ char* customer = std::strstr(tmp_string, "Customer");
+ char* recommends = std::strstr(tmp_string, "Recommends");
+ char* complaints = std::strstr(tmp_string, "Complaints");
+ if (customer) {
+ ASSERT_TRUE((recommends != nullptr) ^ (complaints != nullptr));
+ if (recommends) good_count++;
+ if (complaints) bad_count++;
+ }
+ }
+}
+
+TEST(TpchNode, ScaleFactor) {
+ ExecContext ctx(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(&ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), 0.25f);
+ ExecNode* table = *gen.Supplier();
+ AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+ Declaration sink("sink", {Declaration::Input(table)}, SinkNodeOptions{&sink_gen});
+ std::ignore = *sink.AddToPlan(plan.get());
Review comment:
Please don't ignore it but `ASSERT_OK` on it.
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,631 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <regex>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+namespace internal {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+// Verifies that the data is valid Arrow and ensures it's not null.
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values) {
+ ASSERT_EQ(d.array()->buffers[0].get(), nullptr);
+ ASSERT_OK(d.make_array()->ValidateFull());
+ }
+}
+
+// Verifies that each element is seen exactly once, and that it's between min and max
+// inclusive
+void VerifyUniqueKey(std::unordered_set<int32_t>* seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen->insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const util::string_view& row, const char* prefix,
+ const int64_t i, const int32_t* nums,
+ bool verify_padding) {
+ size_t num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_TRUE(row.starts_with(prefix)) << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row.data() + num_offset;
+ int64_t num = 0;
+ size_t ibyte = num_offset;
+ // Parse the number out
+ for (; *num_str && ibyte < row.size(); ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ // If nums is not null, ensure it matches the parsed number
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ // TPC-H requires only ever requires padding up to 9 digits, so we ensure that
+ // the total length of the string was at least 9 (could be more for bigger numbers).
+ if (verify_padding) {
+ int64_t num_chars = static_cast<int64_t>(ibyte - num_offset);
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+// Verifies that each row is the string "prefix" followed by a number. If numbers is not
+// EMPTY, it also checks that the number following the prefix is equal to the
+// corresponding row in numbers. Some TPC-H data is padded to 9 zeros, which this function
+// can optionally verify as well. This string function verifies fixed width columns.
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ util::string_view view(row, byte_width);
+ VerifyStringAndNumber_Single(view, prefix, i, nums, verify_padding);
+ }
+}
+
+// Same as above but for variable length columns
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ util::string_view view(str + start, str_len);
+ VerifyStringAndNumber_Single(view, prefix, i, nums, verify_padding);
+ }
+}
+
+// Verifies that each row is a V-string, which is defined in the spec to be
+// a string of random length between min_length and max_length, that is composed
+// of alphanumeric characters, commas, or spaces.
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = off[i];
+ int32_t end = off[i + 1];
+ int32_t str_len = end - start;
+ ASSERT_LE(str_len, max_length);
+ ASSERT_GE(str_len, min_length);
+ for (int32_t i = start; i < end; i++) {
+ bool is_valid =
+ std::isdigit(str[i]) || std::isalpha(str[i]) || str[i] == ',' || str[i] == ' ';
+ ASSERT_TRUE(is_valid) << "Character " << str[i]
+ << " is not a digit, a letter, a comma, or a space";
+ }
+ }
+}
+
+// Verifies that each 32-bit element modulo "mod" is between min and max.
+void VerifyModuloBetween(const Datum& d, int32_t min, int32_t max, int32_t mod) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t m = n[i] % mod;
+ ASSERT_GE(m, min) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ ASSERT_LE(m, max) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ }
+}
+
+// Verifies that each 32-bit element is between min and max.
+void VerifyAllBetween(const Datum& d, int32_t min, int32_t max) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ ASSERT_GE(n[i], min) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ ASSERT_LE(n[i], max) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ }
+}
+
+void VerifyNationKey(const Datum& d) { VerifyAllBetween(d, 0, 24); }
+
+// Verifies that each row satisfies the phone number spec.
+void VerifyPhone(const Datum& d) {
+ int64_t length = d.length();
+ const char* phones = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ constexpr int kByteWidth = 15; // This is common for all PHONE columns
+ std::regex exp("\\d{2}-\\d{3}-\\d{3}-\\d{4}");
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = phones + i * kByteWidth;
+ ASSERT_TRUE(std::regex_match(row, row + kByteWidth, exp));
+ }
+}
+
+// Verifies that each decimal is between min and max
+void VerifyDecimalsBetween(const Datum& d, int64_t min, int64_t max) {
+ int64_t length = d.length();
+ const Decimal128* decs =
+ reinterpret_cast<const Decimal128*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int64_t val = static_cast<int64_t>(decs[i]);
+ ASSERT_LE(val, max);
+ ASSERT_GE(val, min);
+ }
+}
+
+// Verifies that each variable-length row is a series of words separated by
+// spaces. Number of words is determined by the number of spaces.
+void VerifyCorrectNumberOfWords_Varlen(const Datum& d, int num_words) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+
+ int32_t start = offsets[i];
+ int32_t end = offsets[i + 1];
+ int32_t str_len = end - start;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, str + start, str_len);
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = offsets[i]; j < offsets[i + 1]; j++) {
+ bool is_space = str[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(str[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << tmp_str;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << tmp_str;
+ }
+}
+
+// Same as above but for fixed width columns.
+void VerifyCorrectNumberOfWords_FixedWidth(const Datum& d, int num_words,
+ int byte_width) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+ const char* row = str + i * byte_width;
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = 0; j < byte_width && row[j]; j++) {
+ bool is_space = row[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(row[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << row;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << row;
+ }
+}
+
+// Verifies that each row of the single-byte-wide column is one of the possibilities.
+void VerifyOneOf(const Datum& d, const std::unordered_set<char>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++)
+ ASSERT_TRUE(possibilities.find(col[i]) != possibilities.end());
+}
+
+// Verifies that each fixed-width row is one of the possibilities
+void VerifyOneOf(const Datum& d, int32_t byte_width,
+ const std::unordered_set<util::string_view>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = col + i * byte_width;
+ int32_t row_len = 0;
+ while (row[row_len] && row_len < byte_width) row_len++;
+ util::string_view view(row, row_len);
+ ASSERT_TRUE(possibilities.find(view) != possibilities.end())
+ << view << " is not a valid string.";
+ }
+}
+
+// Counts the number of instances of each integer
+void CountInstances(std::unordered_map<int32_t, int32_t>& counts, const Datum& d) {
+ int64_t length = d.length();
+ const int32_t* nums = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) counts[nums[i]]++;
+}
+
+// For the S_COMMENT column, some of the columns must be modified to contain
+// "Customer...Complaints" or "Customer...Recommends". This function counts the number of
+// good and bad comments.
+void CountModifiedComments(const Datum& d, int& good_count, int& bad_count) {
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ // Length of S_COMMENT is at most 100
+ char tmp_string[101];
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + offsets[i];
+ int32_t row_length = offsets[i + 1] - offsets[i];
+ std::memset(tmp_string, 0, sizeof(tmp_string));
+ std::memcpy(tmp_string, row, row_length);
+ char* customer = std::strstr(tmp_string, "Customer");
+ char* recommends = std::strstr(tmp_string, "Recommends");
+ char* complaints = std::strstr(tmp_string, "Complaints");
+ if (customer) {
+ ASSERT_TRUE((recommends != nullptr) ^ (complaints != nullptr));
+ if (recommends) good_count++;
+ if (complaints) bad_count++;
+ }
+ }
+}
+
+TEST(TpchNode, ScaleFactor) {
+ ExecContext ctx(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(&ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), 0.25f);
Review comment:
Looks like there's some shared setup code in all those tests that could be factored out into a test fixture?
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,631 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <regex>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+namespace internal {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+// Verifies that the data is valid Arrow and ensures it's not null.
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values) {
+ ASSERT_EQ(d.array()->buffers[0].get(), nullptr);
+ ASSERT_OK(d.make_array()->ValidateFull());
+ }
+}
+
+// Verifies that each element is seen exactly once, and that it's between min and max
+// inclusive
+void VerifyUniqueKey(std::unordered_set<int32_t>* seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen->insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const util::string_view& row, const char* prefix,
+ const int64_t i, const int32_t* nums,
+ bool verify_padding) {
+ size_t num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_TRUE(row.starts_with(prefix)) << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row.data() + num_offset;
+ int64_t num = 0;
+ size_t ibyte = num_offset;
+ // Parse the number out
+ for (; *num_str && ibyte < row.size(); ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ // If nums is not null, ensure it matches the parsed number
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ // TPC-H requires only ever requires padding up to 9 digits, so we ensure that
+ // the total length of the string was at least 9 (could be more for bigger numbers).
+ if (verify_padding) {
+ int64_t num_chars = static_cast<int64_t>(ibyte - num_offset);
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+// Verifies that each row is the string "prefix" followed by a number. If numbers is not
+// EMPTY, it also checks that the number following the prefix is equal to the
+// corresponding row in numbers. Some TPC-H data is padded to 9 zeros, which this function
+// can optionally verify as well. This string function verifies fixed width columns.
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ util::string_view view(row, byte_width);
+ VerifyStringAndNumber_Single(view, prefix, i, nums, verify_padding);
+ }
+}
+
+// Same as above but for variable length columns
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ util::string_view view(str + start, str_len);
+ VerifyStringAndNumber_Single(view, prefix, i, nums, verify_padding);
+ }
+}
+
+// Verifies that each row is a V-string, which is defined in the spec to be
+// a string of random length between min_length and max_length, that is composed
+// of alphanumeric characters, commas, or spaces.
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = off[i];
+ int32_t end = off[i + 1];
+ int32_t str_len = end - start;
+ ASSERT_LE(str_len, max_length);
+ ASSERT_GE(str_len, min_length);
+ for (int32_t i = start; i < end; i++) {
+ bool is_valid =
+ std::isdigit(str[i]) || std::isalpha(str[i]) || str[i] == ',' || str[i] == ' ';
+ ASSERT_TRUE(is_valid) << "Character " << str[i]
+ << " is not a digit, a letter, a comma, or a space";
+ }
+ }
+}
+
+// Verifies that each 32-bit element modulo "mod" is between min and max.
+void VerifyModuloBetween(const Datum& d, int32_t min, int32_t max, int32_t mod) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t m = n[i] % mod;
+ ASSERT_GE(m, min) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ ASSERT_LE(m, max) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ }
+}
+
+// Verifies that each 32-bit element is between min and max.
+void VerifyAllBetween(const Datum& d, int32_t min, int32_t max) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ ASSERT_GE(n[i], min) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ ASSERT_LE(n[i], max) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ }
+}
+
+void VerifyNationKey(const Datum& d) { VerifyAllBetween(d, 0, 24); }
+
+// Verifies that each row satisfies the phone number spec.
+void VerifyPhone(const Datum& d) {
+ int64_t length = d.length();
+ const char* phones = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ constexpr int kByteWidth = 15; // This is common for all PHONE columns
+ std::regex exp("\\d{2}-\\d{3}-\\d{3}-\\d{4}");
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = phones + i * kByteWidth;
+ ASSERT_TRUE(std::regex_match(row, row + kByteWidth, exp));
+ }
+}
+
+// Verifies that each decimal is between min and max
+void VerifyDecimalsBetween(const Datum& d, int64_t min, int64_t max) {
+ int64_t length = d.length();
+ const Decimal128* decs =
+ reinterpret_cast<const Decimal128*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int64_t val = static_cast<int64_t>(decs[i]);
+ ASSERT_LE(val, max);
+ ASSERT_GE(val, min);
+ }
+}
+
+// Verifies that each variable-length row is a series of words separated by
+// spaces. Number of words is determined by the number of spaces.
+void VerifyCorrectNumberOfWords_Varlen(const Datum& d, int num_words) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+
+ int32_t start = offsets[i];
+ int32_t end = offsets[i + 1];
+ int32_t str_len = end - start;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, str + start, str_len);
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = offsets[i]; j < offsets[i + 1]; j++) {
+ bool is_space = str[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(str[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << tmp_str;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << tmp_str;
+ }
+}
+
+// Same as above but for fixed width columns.
+void VerifyCorrectNumberOfWords_FixedWidth(const Datum& d, int num_words,
+ int byte_width) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+ const char* row = str + i * byte_width;
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = 0; j < byte_width && row[j]; j++) {
+ bool is_space = row[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(row[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << row;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << row;
+ }
+}
+
+// Verifies that each row of the single-byte-wide column is one of the possibilities.
+void VerifyOneOf(const Datum& d, const std::unordered_set<char>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++)
+ ASSERT_TRUE(possibilities.find(col[i]) != possibilities.end());
+}
+
+// Verifies that each fixed-width row is one of the possibilities
+void VerifyOneOf(const Datum& d, int32_t byte_width,
+ const std::unordered_set<util::string_view>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = col + i * byte_width;
+ int32_t row_len = 0;
+ while (row[row_len] && row_len < byte_width) row_len++;
+ util::string_view view(row, row_len);
+ ASSERT_TRUE(possibilities.find(view) != possibilities.end())
+ << view << " is not a valid string.";
+ }
+}
+
+// Counts the number of instances of each integer
+void CountInstances(std::unordered_map<int32_t, int32_t>& counts, const Datum& d) {
+ int64_t length = d.length();
+ const int32_t* nums = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) counts[nums[i]]++;
+}
+
+// For the S_COMMENT column, some of the columns must be modified to contain
+// "Customer...Complaints" or "Customer...Recommends". This function counts the number of
+// good and bad comments.
+void CountModifiedComments(const Datum& d, int& good_count, int& bad_count) {
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ // Length of S_COMMENT is at most 100
+ char tmp_string[101];
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + offsets[i];
+ int32_t row_length = offsets[i + 1] - offsets[i];
+ std::memset(tmp_string, 0, sizeof(tmp_string));
+ std::memcpy(tmp_string, row, row_length);
+ char* customer = std::strstr(tmp_string, "Customer");
+ char* recommends = std::strstr(tmp_string, "Recommends");
+ char* complaints = std::strstr(tmp_string, "Complaints");
+ if (customer) {
+ ASSERT_TRUE((recommends != nullptr) ^ (complaints != nullptr));
+ if (recommends) good_count++;
+ if (complaints) bad_count++;
+ }
+ }
+}
+
+TEST(TpchNode, ScaleFactor) {
+ ExecContext ctx(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(&ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), 0.25f);
+ ExecNode* table = *gen.Supplier();
+ AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+ Declaration sink("sink", {Declaration::Input(table)}, SinkNodeOptions{&sink_gen});
+ std::ignore = *sink.AddToPlan(plan.get());
+ auto fut = StartAndCollect(plan.get(), sink_gen);
+ auto res = *fut.MoveResult();
+
+ int64_t kExpectedRows = 2500;
+ int64_t num_rows = 0;
+ for (auto& batch : res) num_rows += batch.length;
+ ASSERT_EQ(num_rows, kExpectedRows);
+}
+
+TEST(TpchNode, Supplier) {
+ ExecContext ctx(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(&ctx);
+ TpchGen gen = *TpchGen::Make(plan.get());
+ ExecNode* table = *gen.Supplier();
Review comment:
Here is as well, check for unexpected errors instead of potentially crashing the test process...
```suggestion
ASSERT_OK_AND_ASSIGN(std::shared_ptr<ExecPlan> plan, ExecPlan::Make(&ctx));
ASSERT_OK_AND_ASSIGN(TpchGen gen, TpchGen::Make(plan.get(), 0.25f));
```
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,123 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <benchmark/benchmark.h>
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+namespace internal {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<double>(scale_factor));
+
+ ExecNode* lineitem =
+ *gen.Lineitem({"L_QUANTITY", "L_EXTENDEDPRICE", "L_TAX", "L_DISCOUNT", "L_SHIPDATE",
+ "L_RETURNFLAG", "L_LINESTATUS"});
+
+ auto sept_2_1998 = std::make_shared<Date32Scalar>(
+ 10471); // September 2, 1998 is 10471 days after January 1, 1970
+ Expression filter =
+ less_equal(field_ref("L_SHIPDATE"), literal(std::move(sept_2_1998)));
+ FilterNodeOptions filter_opts(filter);
+
+ Expression l_returnflag = field_ref("L_RETURNFLAG");
+ Expression l_linestatus = field_ref("L_LINESTATUS");
+ Expression quantity = field_ref("L_QUANTITY");
+ Expression base_price = field_ref("L_EXTENDEDPRICE");
+
+ std::shared_ptr<Decimal128Scalar> decimal_1 =
+ std::make_shared<Decimal128Scalar>(Decimal128{0, 100}, decimal(12, 2));
+ Expression discount_multiplier =
+ call("subtract", {literal(decimal_1), field_ref("L_DISCOUNT")});
+ Expression tax_multiplier = call("add", {literal(decimal_1), field_ref("L_TAX")});
+ Expression disc_price =
+ call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier});
+ Expression charge =
+ call("multiply",
+ {call("cast",
+ {call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier})},
+ compute::CastOptions::Unsafe(decimal(12, 2))),
+ tax_multiplier});
+ Expression discount = field_ref("L_DISCOUNT");
+
+ std::vector<Expression> projection_list = {l_returnflag, l_linestatus, quantity,
+ base_price, disc_price, charge,
+ quantity, base_price, discount};
+ std::vector<std::string> project_names = {
+ "l_returnflag", "l_linestatus", "sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price", "avg_disc"};
+ ProjectNodeOptions project_opts(std::move(projection_list), std::move(project_names));
+
+ ScalarAggregateOptions sum_opts = ScalarAggregateOptions::Defaults();
+ CountOptions count_opts(CountOptions::CountMode::ALL);
+ std::vector<arrow::compute::internal::Aggregate> aggs = {
+ {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts},
+ {"hash_sum", &sum_opts}, {"hash_mean", &sum_opts}, {"hash_mean", &sum_opts},
+ {"hash_mean", &sum_opts}, {"hash_count", &count_opts}};
+
+ std::vector<FieldRef> to_aggregate = {"sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price",
+ "avg_disc", "sum_qty"};
+
+ std::vector<std::string> names = {"sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price",
+ "avg_disc", "count_order"};
+
+ std::vector<FieldRef> keys = {"l_returnflag", "l_linestatus"};
+ AggregateNodeOptions agg_opts(aggs, to_aggregate, names, keys);
+
+ SortKey l_returnflag_key("l_returnflag");
+ SortKey l_linestatus_key("l_linestatus");
+ SortOptions sort_opts({l_returnflag_key, l_linestatus_key});
+ OrderBySinkNodeOptions order_by_opts(sort_opts, sink_gen);
+
+ Declaration filter_decl("filter", {Declaration::Input(lineitem)}, filter_opts);
+ Declaration project_decl("project", project_opts);
+ Declaration aggregate_decl("aggregate", agg_opts);
+ Declaration orderby_decl("order_by_sink", order_by_opts);
+
+ Declaration q1 =
+ Declaration::Sequence({filter_decl, project_decl, aggregate_decl, orderby_decl});
+ std::ignore = *q1.AddToPlan(plan.get());
+ return plan;
+}
+
+static void BM_Tpch_Q1(benchmark::State& st) {
+ for (auto _ : st) {
+ st.PauseTiming();
+ AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+ std::shared_ptr<ExecPlan> plan = Plan_Q1(&sink_gen, static_cast<int>(st.range(0)));
+ st.ResumeTiming();
+ auto fut = StartAndCollect(plan.get(), sink_gen);
+ auto res = *fut.MoveResult();
+ }
+}
Review comment:
I see this benchmark only outputs a timing in nanoseconds. Do we want to output a more meaningful indicator as well (items per second? something else?)?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082215097
@save-buffer I would have appreciated if you asked for the changeset motivation instead of bluntly reverting it. The point here was to minimize the exposure of implementation details in a `.h` file, and also consequently minimize the set of transitive inclusions (as a common concern with C++ compilation times).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] ursabot edited a comment on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
ursabot edited a comment on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082832095
Benchmark runs are scheduled for baseline = 0deefd1838bf4706232b90b72da010698010f7b6 and contender = 50fab73da58c9f76fbe675d661e56087e2f5fb52. 50fab73da58c9f76fbe675d661e56087e2f5fb52 is a master commit associated with this PR. Results will be available as each benchmark for each run completes.
Conbench compare runs links:
[Finished :arrow_down:0.0% :arrow_up:0.0%] [ec2-t3-xlarge-us-east-2](https://conbench.ursa.dev/compare/runs/3f8ac59fddef426cac31b6b2826ce5a6...e5b78c70b6fe4a699c4cb8b967fe5736/)
[Scheduled] [test-mac-arm](https://conbench.ursa.dev/compare/runs/02b74ab900d440f8b5361300f15ad15c...b870e78b617e41e98c2ade7241e0aa06/)
[Scheduled] [ursa-i9-9960x](https://conbench.ursa.dev/compare/runs/1f7740546747452aa61fece1295aa0e3...f66d8edd43724d9988628cbe4a4e4d50/)
[Scheduled] [ursa-thinkcentre-m75q](https://conbench.ursa.dev/compare/runs/c3c55d6999fe4d7eabec73cbdfa1fad6...4b9f45f6442447d3a9d758598d0e14d3/)
Buildkite builds:
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/415| `50fab73d` ec2-t3-xlarge-us-east-2>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/414| `0deefd18` ec2-t3-xlarge-us-east-2>
Supported benchmarks:
ec2-t3-xlarge-us-east-2: Supported benchmark langs: Python, R. Runs only benchmarks with cloud = True
test-mac-arm: Supported benchmark langs: C++, Python, R
ursa-i9-9960x: Supported benchmark langs: Python, R, JavaScript
ursa-thinkcentre-m75q: Supported benchmark langs: C++, Java
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r837743650
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -42,40 +42,23 @@ class ARROW_EXPORT TpchGen {
* table from that single TpchGen instance. Note: Every batch will be scheduled as a new
* task using the ExecPlan's scheduler.
*/
- static Result<TpchGen> Make(ExecPlan* plan, double scale_factor = 1.0,
- int64_t batch_size = 4096,
- util::optional<int64_t> seed = util::nullopt);
+ static Result<std::unique_ptr<TpchGen>> Make(
Review comment:
This does not need to be a unique_ptr. The whole point of using shared_ptr for the order/lineitem and part/partsupp generators is so that this TpchGen object can safely die.
Also none of the methods in this class need to be virtual.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] ursabot edited a comment on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
ursabot edited a comment on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082832095
Benchmark runs are scheduled for baseline = 0deefd1838bf4706232b90b72da010698010f7b6 and contender = 50fab73da58c9f76fbe675d661e56087e2f5fb52. 50fab73da58c9f76fbe675d661e56087e2f5fb52 is a master commit associated with this PR. Results will be available as each benchmark for each run completes.
Conbench compare runs links:
[Finished :arrow_down:0.0% :arrow_up:0.0%] [ec2-t3-xlarge-us-east-2](https://conbench.ursa.dev/compare/runs/3f8ac59fddef426cac31b6b2826ce5a6...e5b78c70b6fe4a699c4cb8b967fe5736/)
[Scheduled] [test-mac-arm](https://conbench.ursa.dev/compare/runs/02b74ab900d440f8b5361300f15ad15c...b870e78b617e41e98c2ade7241e0aa06/)
[Scheduled] [ursa-i9-9960x](https://conbench.ursa.dev/compare/runs/1f7740546747452aa61fece1295aa0e3...f66d8edd43724d9988628cbe4a4e4d50/)
[Scheduled] [ursa-thinkcentre-m75q](https://conbench.ursa.dev/compare/runs/c3c55d6999fe4d7eabec73cbdfa1fad6...4b9f45f6442447d3a9d758598d0e14d3/)
Supported benchmarks:
ec2-t3-xlarge-us-east-2: Supported benchmark langs: Python, R. Runs only benchmarks with cloud = True
test-mac-arm: Supported benchmark langs: C++, Python, R
ursa-i9-9960x: Supported benchmark langs: Python, R, JavaScript
ursa-thinkcentre-m75q: Supported benchmark langs: C++, Java
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] ursabot edited a comment on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
ursabot edited a comment on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082832095
Benchmark runs are scheduled for baseline = 0deefd1838bf4706232b90b72da010698010f7b6 and contender = 50fab73da58c9f76fbe675d661e56087e2f5fb52. 50fab73da58c9f76fbe675d661e56087e2f5fb52 is a master commit associated with this PR. Results will be available as each benchmark for each run completes.
Conbench compare runs links:
[Finished :arrow_down:0.0% :arrow_up:0.0%] [ec2-t3-xlarge-us-east-2](https://conbench.ursa.dev/compare/runs/3f8ac59fddef426cac31b6b2826ce5a6...e5b78c70b6fe4a699c4cb8b967fe5736/)
[Finished :arrow_down:0.38% :arrow_up:0.33%] [test-mac-arm](https://conbench.ursa.dev/compare/runs/02b74ab900d440f8b5361300f15ad15c...b870e78b617e41e98c2ade7241e0aa06/)
[Scheduled] [ursa-i9-9960x](https://conbench.ursa.dev/compare/runs/1f7740546747452aa61fece1295aa0e3...f66d8edd43724d9988628cbe4a4e4d50/)
[Finished :arrow_down:0.04% :arrow_up:0.26%] [ursa-thinkcentre-m75q](https://conbench.ursa.dev/compare/runs/c3c55d6999fe4d7eabec73cbdfa1fad6...4b9f45f6442447d3a9d758598d0e14d3/)
Buildkite builds:
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/415| `50fab73d` ec2-t3-xlarge-us-east-2>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-test-mac-arm/builds/401| `50fab73d` test-mac-arm>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-thinkcentre-m75q/builds/411| `50fab73d` ursa-thinkcentre-m75q>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/414| `0deefd18` ec2-t3-xlarge-us-east-2>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-test-mac-arm/builds/400| `0deefd18` test-mac-arm>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-thinkcentre-m75q/builds/410| `0deefd18` ursa-thinkcentre-m75q>
Supported benchmarks:
ec2-t3-xlarge-us-east-2: Supported benchmark langs: Python, R. Runs only benchmarks with cloud = True
test-mac-arm: Supported benchmark langs: C++, Python, R
ursa-i9-9960x: Supported benchmark langs: Python, R, JavaScript
ursa-thinkcentre-m75q: Supported benchmark langs: C++, Java
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] ursabot edited a comment on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
ursabot edited a comment on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082832095
Benchmark runs are scheduled for baseline = 0deefd1838bf4706232b90b72da010698010f7b6 and contender = 50fab73da58c9f76fbe675d661e56087e2f5fb52. 50fab73da58c9f76fbe675d661e56087e2f5fb52 is a master commit associated with this PR. Results will be available as each benchmark for each run completes.
Conbench compare runs links:
[Finished :arrow_down:0.0% :arrow_up:0.0%] [ec2-t3-xlarge-us-east-2](https://conbench.ursa.dev/compare/runs/3f8ac59fddef426cac31b6b2826ce5a6...e5b78c70b6fe4a699c4cb8b967fe5736/)
[Scheduled] [test-mac-arm](https://conbench.ursa.dev/compare/runs/02b74ab900d440f8b5361300f15ad15c...b870e78b617e41e98c2ade7241e0aa06/)
[Scheduled] [ursa-i9-9960x](https://conbench.ursa.dev/compare/runs/1f7740546747452aa61fece1295aa0e3...f66d8edd43724d9988628cbe4a4e4d50/)
[Scheduled] [ursa-thinkcentre-m75q](https://conbench.ursa.dev/compare/runs/c3c55d6999fe4d7eabec73cbdfa1fad6...4b9f45f6442447d3a9d758598d0e14d3/)
Buildkite builds:
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/415| `50fab73d` ec2-t3-xlarge-us-east-2>
[Finished] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ec2-t3-xlarge-us-east-2/builds/414| `0deefd18` ec2-t3-xlarge-us-east-2>
[Scheduled] <https://buildkite.com/apache-arrow/arrow-bci-benchmark-on-ursa-thinkcentre-m75q/builds/410| `0deefd18` ursa-thinkcentre-m75q>
Supported benchmarks:
ec2-t3-xlarge-us-east-2: Supported benchmark langs: Python, R. Runs only benchmarks with cloud = True
test-mac-arm: Supported benchmark langs: C++, Python, R
ursa-i9-9960x: Supported benchmark langs: Python, R, JavaScript
ursa-thinkcentre-m75q: Supported benchmark langs: C++, Java
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829478094
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
Review comment:
I agree we want users to use this but I don't think that means we will need the node in a public namespace. Eventually, we will add function registry and Substrait integrations. At that point we will probably create some kind of "options" object which will be in the public API. At that point the nodes could even be in anonymous namespace, however I think the direct construction is necessary now and will probably still be useful in the future for unit testing. I think `arrow::compute::internal` is reasonable.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829708111
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
Review comment:
i made it take a seed, defaulting to generating a nondeterministic one.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829708651
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row + num_offset;
+ int64_t num = 0;
+ int ibyte = static_cast<int>(num_offset);
+ for (; *num_str && ibyte < byte_width; ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ if (verify_padding) {
+ int num_chars = ibyte - num_offset;
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ VerifyStringAndNumber_Single(row, prefix, i, nums, byte_width, verify_padding);
+ }
+}
+
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ char tmp_str[256] = {};
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ std::memcpy(tmp_str, str + start, str_len);
+ VerifyStringAndNumber_Single(tmp_str, prefix, i, nums, sizeof(tmp_str),
+ verify_padding);
+ }
Review comment:
i made it verify that there are no nulls when we verify the data.
regarding visitors, they look clunky and the code that implements it isn't very readable. i'd prefer to stick with the simple solution.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829775306
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<float>(scale_factor));
+
+ ExecNode* lineitem =
+ *gen.Lineitem({"L_QUANTITY", "L_EXTENDEDPRICE", "L_TAX", "L_DISCOUNT", "L_SHIPDATE",
+ "L_RETURNFLAG", "L_LINESTATUS"});
+
+ std::shared_ptr<Date32Scalar> sept_2_1998 = std::make_shared<Date32Scalar>(
+ 10471); // September 2, 1998 is 10471 days after January 1, 1970
+ Expression filter =
+ less_equal(field_ref("L_SHIPDATE"), literal(std::move(sept_2_1998)));
+ FilterNodeOptions filter_opts(filter);
+
+ Expression l_returnflag = field_ref("L_RETURNFLAG");
+ Expression l_linestatus = field_ref("L_LINESTATUS");
+ Expression quantity = field_ref("L_QUANTITY");
+ Expression base_price = field_ref("L_EXTENDEDPRICE");
+
+ std::shared_ptr<Decimal128Scalar> decimal_1 =
+ std::make_shared<Decimal128Scalar>(Decimal128{0, 100}, decimal(12, 2));
+ Expression discount_multiplier =
+ call("subtract", {literal(decimal_1), field_ref("L_DISCOUNT")});
+ Expression tax_multiplier = call("add", {literal(decimal_1), field_ref("L_TAX")});
+ Expression disc_price =
+ call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier});
+ Expression charge =
+ call("multiply",
+ {call("cast",
+ {call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier})},
+ compute::CastOptions::Unsafe(decimal(12, 2))),
+ tax_multiplier});
+ Expression discount = field_ref("L_DISCOUNT");
+
+ std::vector<Expression> projection_list = {l_returnflag, l_linestatus, quantity,
+ base_price, disc_price, charge,
+ quantity, base_price, discount};
+ std::vector<std::string> project_names = {
+ "l_returnflag", "l_linestatus", "sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price", "avg_disc"};
+ ProjectNodeOptions project_opts(std::move(projection_list));
+
+ ScalarAggregateOptions sum_opts = ScalarAggregateOptions::Defaults();
+ CountOptions count_opts(CountOptions::CountMode::ALL);
+ std::vector<arrow::compute::internal::Aggregate> aggs = {
+ {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts},
+ {"hash_sum", &sum_opts}, {"hash_mean", &sum_opts}, {"hash_mean", &sum_opts},
+ {"hash_mean", &sum_opts}, {"hash_count", &count_opts}};
+
+ std::vector<FieldRef> cols = {2, 3, 4, 5, 6, 7, 8, 2};
Review comment:
Because one doesn't need an indirection to know what these columns are.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1068020835
> Thanks for addressing my comments. A few remaining nits but I'm content with the c++ part. @jonkeane do you need me to take a look at the R?
If you wouldn't mind, that would be fabulous (especially the C++ in the r directory...) I've still not been able to get the `tpch_dbgen_write` function to work (even with my hard coded version here) — but I haven't been able to track down where it's going wrong. I had put in a few messages to try, and it appears that it's getting to producing, but never completes and I get segfaults:
```
> arrow:::tpch_dbgen_write("lineitem", scale_factor = 1, path = "~/Desktop")
Just after consume
Just after validate
Just after start
*** caught segfault ***
*** caught segfault ***
address 0x0, cause 'memory not mapped'
*** caught segfault ***
*** caught segfault ***
*** caught segfault ***
address 0x0, cause 'memory not mapped'
address 0x0, cause 'memory not mapped'
address 0x0, cause 'memory not mapped'
address 0x0, cause 'memory not mapped'
*** caught segfault ***
address 0x0, cause 'memory not mapped'
*** caught segfault ***
address 0x0, cause 'memory not mapped'
*** caught segfault ***
address 0x0, cause 'memory not mapped'
```
I've also attached [a crash report](https://github.com/apache/arrow/files/8253958/tpch_write_crash.txt)
We can move those out of this PR if we want to get this done and merge it separately (though I'll probably still need some C++ help to figure out what's going on with that if we do it separately)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1056020026
Good point about tests - I will add some.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1057572921
Well that might explain the seg fault I got! Thanks for that I’ll try piping that in.
I have some thoughts about if that actually works, but I need to dig some more. My issues with trying to use it that way preciously might have been from the R bindings and not C++ so I might be able to work around them. I know the single file write is not standard, and unlikely to be the right thing for most people, but in this case it helps a lot with compatibility.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829330091
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
+class OrdersAndLineItemGenerator;
+class PartAndPartSupplierGenerator;
+
+class ARROW_EXPORT TpchGen {
+ public:
+ /*
+ * \brief Create a factory for nodes that generate TPC-H data
+ *
+ * Note: Individual tables will reference each other. It is important that you only
+ * create a single TpchGen instance for each plan and then you can create nodes for each
+ * table from that single TpchGen instance. Note: Every batch will be scheduled as a new
+ * task using the ExecPlan's scheduler.
+ */
+ static Result<TpchGen> Make(ExecPlan* plan, float scale_factor = 1.0f,
+ int64_t batch_size = 4096);
+
+ Result<ExecNode*> Supplier(std::vector<std::string> columns = {});
Review comment:
done
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
Review comment:
done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou edited a comment on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou edited a comment on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1082215097
@save-buffer I would have appreciated if you had asked for the changeset motivation instead of bluntly reverting it. The point here was to minimize the exposure of implementation details in a `.h` file, and also consequently minimize the set of transitive inclusions (as a common concern with C++ compilation times).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou closed pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou closed pull request #12537:
URL: https://github.com/apache/arrow/pull/12537
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829181919
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
Review comment:
Nit: use bracket around external dependencies.
```suggestion
#include <benchmark/benchmark.h>
```
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<float>(scale_factor));
+
+ ExecNode* lineitem =
+ *gen.Lineitem({"L_QUANTITY", "L_EXTENDEDPRICE", "L_TAX", "L_DISCOUNT", "L_SHIPDATE",
+ "L_RETURNFLAG", "L_LINESTATUS"});
+
+ std::shared_ptr<Date32Scalar> sept_2_1998 = std::make_shared<Date32Scalar>(
+ 10471); // September 2, 1998 is 10471 days after January 1, 1970
+ Expression filter =
+ less_equal(field_ref("L_SHIPDATE"), literal(std::move(sept_2_1998)));
+ FilterNodeOptions filter_opts(filter);
+
+ Expression l_returnflag = field_ref("L_RETURNFLAG");
+ Expression l_linestatus = field_ref("L_LINESTATUS");
+ Expression quantity = field_ref("L_QUANTITY");
+ Expression base_price = field_ref("L_EXTENDEDPRICE");
+
+ std::shared_ptr<Decimal128Scalar> decimal_1 =
+ std::make_shared<Decimal128Scalar>(Decimal128{0, 100}, decimal(12, 2));
+ Expression discount_multiplier =
+ call("subtract", {literal(decimal_1), field_ref("L_DISCOUNT")});
+ Expression tax_multiplier = call("add", {literal(decimal_1), field_ref("L_TAX")});
+ Expression disc_price =
+ call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier});
+ Expression charge =
+ call("multiply",
+ {call("cast",
+ {call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier})},
+ compute::CastOptions::Unsafe(decimal(12, 2))),
+ tax_multiplier});
+ Expression discount = field_ref("L_DISCOUNT");
+
+ std::vector<Expression> projection_list = {l_returnflag, l_linestatus, quantity,
+ base_price, disc_price, charge,
+ quantity, base_price, discount};
+ std::vector<std::string> project_names = {
+ "l_returnflag", "l_linestatus", "sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price", "avg_disc"};
+ ProjectNodeOptions project_opts(std::move(projection_list));
+
+ ScalarAggregateOptions sum_opts = ScalarAggregateOptions::Defaults();
+ CountOptions count_opts(CountOptions::CountMode::ALL);
+ std::vector<arrow::compute::internal::Aggregate> aggs = {
+ {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts},
+ {"hash_sum", &sum_opts}, {"hash_mean", &sum_opts}, {"hash_mean", &sum_opts},
+ {"hash_mean", &sum_opts}, {"hash_count", &count_opts}};
+
+ std::vector<FieldRef> cols = {2, 3, 4, 5, 6, 7, 8, 2};
+
+ std::vector<std::string> names = {"sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price",
+ "avg_disc", "count_order"};
+
+ std::vector<FieldRef> keys = {"L_RETURNFLAG", "L_LINESTATUS"};
+ AggregateNodeOptions agg_opts(aggs, cols, names, keys);
+
+ SortKey l_returnflag_key("L_RETURNFLAG");
+ SortKey l_linestatus_key("L_LINESTATUS");
+ SortOptions sort_opts({l_returnflag_key, l_linestatus_key});
+ OrderBySinkNodeOptions order_by_opts(sort_opts, sink_gen);
+
+ Declaration filter_decl("filter", {Declaration::Input(lineitem)}, filter_opts);
+ Declaration project_decl("project", project_opts);
+ Declaration aggregate_decl("aggregate", agg_opts);
+ Declaration orderby_decl("order_by_sink", order_by_opts);
+
+ Declaration q1 =
+ Declaration::Sequence({filter_decl, project_decl, aggregate_decl, orderby_decl});
+ std::ignore = *q1.AddToPlan(plan.get());
+ return plan;
+}
+
+static void BM_Tpch_Q1(benchmark::State& st) {
+ for (auto _ : st) {
+ st.PauseTiming();
+ AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+ std::shared_ptr<ExecPlan> plan = Plan_Q1(&sink_gen, static_cast<int>(st.range(0)));
+ st.ResumeTiming();
+ auto fut = StartAndCollect(plan.get(), sink_gen);
+ auto res = *fut.MoveResult();
+ }
+}
+
+BENCHMARK(BM_Tpch_Q1)->Args({1})->ArgNames({"SF"});
Review comment:
"SF" is for scale factor? Wouldn't it be nicer to spell it out in full?
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
+class OrdersAndLineItemGenerator;
+class PartAndPartSupplierGenerator;
+
+class ARROW_EXPORT TpchGen {
+ public:
+ /*
+ * \brief Create a factory for nodes that generate TPC-H data
+ *
+ * Note: Individual tables will reference each other. It is important that you only
+ * create a single TpchGen instance for each plan and then you can create nodes for each
+ * table from that single TpchGen instance. Note: Every batch will be scheduled as a new
+ * task using the ExecPlan's scheduler.
+ */
+ static Result<TpchGen> Make(ExecPlan* plan, float scale_factor = 1.0f,
+ int64_t batch_size = 4096);
Review comment:
Is it me, or does 4096 seem a bit small?
##########
File path: cpp/src/arrow/compute/kernels/vector_sort.cc
##########
@@ -893,7 +893,8 @@ class TableSorter {
TableSorter(ExecContext* ctx, uint64_t* indices_begin, uint64_t* indices_end,
const Table& table, const SortOptions& options)
- : ctx_(ctx),
+ : status_(),
+ ctx_(ctx),
Review comment:
Nit, but there's no need to explicitly spell out the default initialization for `status_`.
```suggestion
: ctx_(ctx),
```
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
Review comment:
Can you add a short comment at the beginning of each `Verify...` function explaining what it does?
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
Review comment:
Why aren't you using `util::string_view` or `std::string`?
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<float>(scale_factor));
+
+ ExecNode* lineitem =
+ *gen.Lineitem({"L_QUANTITY", "L_EXTENDEDPRICE", "L_TAX", "L_DISCOUNT", "L_SHIPDATE",
+ "L_RETURNFLAG", "L_LINESTATUS"});
+
+ std::shared_ptr<Date32Scalar> sept_2_1998 = std::make_shared<Date32Scalar>(
+ 10471); // September 2, 1998 is 10471 days after January 1, 1970
+ Expression filter =
+ less_equal(field_ref("L_SHIPDATE"), literal(std::move(sept_2_1998)));
+ FilterNodeOptions filter_opts(filter);
+
+ Expression l_returnflag = field_ref("L_RETURNFLAG");
+ Expression l_linestatus = field_ref("L_LINESTATUS");
+ Expression quantity = field_ref("L_QUANTITY");
+ Expression base_price = field_ref("L_EXTENDEDPRICE");
+
+ std::shared_ptr<Decimal128Scalar> decimal_1 =
+ std::make_shared<Decimal128Scalar>(Decimal128{0, 100}, decimal(12, 2));
+ Expression discount_multiplier =
+ call("subtract", {literal(decimal_1), field_ref("L_DISCOUNT")});
+ Expression tax_multiplier = call("add", {literal(decimal_1), field_ref("L_TAX")});
+ Expression disc_price =
+ call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier});
+ Expression charge =
+ call("multiply",
+ {call("cast",
+ {call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier})},
+ compute::CastOptions::Unsafe(decimal(12, 2))),
+ tax_multiplier});
+ Expression discount = field_ref("L_DISCOUNT");
+
+ std::vector<Expression> projection_list = {l_returnflag, l_linestatus, quantity,
+ base_price, disc_price, charge,
+ quantity, base_price, discount};
+ std::vector<std::string> project_names = {
+ "l_returnflag", "l_linestatus", "sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price", "avg_disc"};
+ ProjectNodeOptions project_opts(std::move(projection_list));
+
+ ScalarAggregateOptions sum_opts = ScalarAggregateOptions::Defaults();
+ CountOptions count_opts(CountOptions::CountMode::ALL);
+ std::vector<arrow::compute::internal::Aggregate> aggs = {
+ {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts},
+ {"hash_sum", &sum_opts}, {"hash_mean", &sum_opts}, {"hash_mean", &sum_opts},
+ {"hash_mean", &sum_opts}, {"hash_count", &count_opts}};
+
+ std::vector<FieldRef> cols = {2, 3, 4, 5, 6, 7, 8, 2};
Review comment:
Without changing anything functionally, perhaps it would be nicer to initialize this vector from field names rather than field indices?
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<float>(scale_factor));
Review comment:
Why `float`? This seems unimportant, but should probably pass a `double` for future-proofing.
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row + num_offset;
+ int64_t num = 0;
+ int ibyte = static_cast<int>(num_offset);
+ for (; *num_str && ibyte < byte_width; ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ if (verify_padding) {
+ int num_chars = ibyte - num_offset;
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ VerifyStringAndNumber_Single(row, prefix, i, nums, byte_width, verify_padding);
+ }
+}
+
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ char tmp_str[256] = {};
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ std::memcpy(tmp_str, str + start, str_len);
+ VerifyStringAndNumber_Single(tmp_str, prefix, i, nums, sizeof(tmp_str),
+ verify_padding);
+ }
+}
+
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = off[i];
+ int32_t end = off[i + 1];
+ int32_t str_len = end - start;
+ ASSERT_LE(str_len, max_length);
+ ASSERT_GE(str_len, min_length);
+ for (int32_t i = start; i < end; i++) {
+ bool is_valid =
+ std::isdigit(str[i]) || std::isalpha(str[i]) || str[i] == ',' || str[i] == ' ';
+ ASSERT_TRUE(is_valid) << "Character " << str[i]
+ << " is not a digit, a letter, a comma, or a space";
+ }
+ }
+}
+
+void VerifyModuloBetween(const Datum& d, int32_t min, int32_t max, int32_t mod) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t m = n[i] % mod;
+ ASSERT_GE(m, min) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ ASSERT_LE(m, max) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ }
+}
+
+void VerifyAllBetween(const Datum& d, int32_t min, int32_t max) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ ASSERT_GE(n[i], min) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ ASSERT_LE(n[i], max) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ }
+}
+
+void VerifyNationKey(const Datum& d) { VerifyAllBetween(d, 0, 24); }
+
+void VerifyPhone(const Datum& d) {
+ int64_t length = d.length();
+ const char* phones = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ constexpr int kByteWidth = 15; // This is common for all PHONE columns
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = phones + i * kByteWidth;
Review comment:
It's probably allowed to use `std::regex` in tests (I see we already use it in some dataset tests).
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
Review comment:
Why is `pcg_random.h` included here? Generally we should avoid including stuff in `.h` files if it's only needed for the implementations in the corresponding `.cc`.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
Review comment:
Why doesn't this use the existing formatting functionalities in `arrow/util/formatting.h`?
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ if (!partsupp_output_queue_.empty()) {
+ ExecBatch result = std::move(partsupp_output_queue_.front());
+ partsupp_output_queue_.pop();
+ return std::move(result);
+ }
+ }
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+ if (!part_cols_.empty()) {
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch part_batch,
+ ExecBatch::Make(std::move(part_result)));
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ part_output_queue_.emplace(std::move(part_batch));
+ }
+ }
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ for (size_t i = 1; i < partsupp_results.size(); i++)
+ partsupp_output_queue_.emplace(std::move(partsupp_results[i]));
+ }
+ return std::move(partsupp_results[0]);
+ }
+
+ private:
+#define FOR_EACH_PART_COLUMN(F) \
+ F(P_PARTKEY) \
+ F(P_NAME) \
+ F(P_MFGR) \
+ F(P_BRAND) \
+ F(P_TYPE) \
+ F(P_SIZE) \
+ F(P_CONTAINER) \
+ F(P_RETAILPRICE) \
+ F(P_COMMENT)
+
+#define FOR_EACH_PARTSUPP_COLUMN(F) \
+ F(PS_PARTKEY) \
+ F(PS_SUPPKEY) \
+ F(PS_AVAILQTY) \
+ F(PS_SUPPLYCOST) \
+ F(PS_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct PART {
+ enum {
+ FOR_EACH_PART_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct PARTSUPP {
+ enum {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, PART::col},
+ const std::unordered_map<std::string, int> kPartNameMap = {
+ FOR_EACH_PART_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, PARTSUPP::col},
+ const std::unordered_map<std::string, int> kPartsuppNameMap = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kPartGenerators = {FOR_EACH_PART_COLUMN(MAKE_FN_ARRAY)};
+ std::vector<GenerateColumnFn> kPartsuppGenerators = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kPartTypes = {
+ int32(), utf8(), fixed_size_binary(25), fixed_size_binary(10),
+ utf8(), int32(), fixed_size_binary(10), decimal(12, 2),
+ utf8(),
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kPartsuppTypes = {
+ int32(), int32(), int32(), decimal(12, 2), utf8(),
+ };
+
+ Status AllocatePartBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.part[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.part_to_generate * byte_width));
+ ArrayData ad(kPartTypes[column], tld.part_to_generate, {nullptr, std::move(buff)});
+ tld.part[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status P_PARTKEY(size_t thread_index) {
Review comment:
Can you add comments describing and explaining the various generation methods below?
We should be able to understand and maintain this code without knowing the TPC-H spec by heart, ideally :-)
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
+class OrdersAndLineItemGenerator;
+class PartAndPartSupplierGenerator;
+
+class ARROW_EXPORT TpchGen {
+ public:
+ /*
+ * \brief Create a factory for nodes that generate TPC-H data
+ *
+ * Note: Individual tables will reference each other. It is important that you only
+ * create a single TpchGen instance for each plan and then you can create nodes for each
+ * table from that single TpchGen instance. Note: Every batch will be scheduled as a new
+ * task using the ExecPlan's scheduler.
+ */
+ static Result<TpchGen> Make(ExecPlan* plan, float scale_factor = 1.0f,
+ int64_t batch_size = 4096);
+
+ Result<ExecNode*> Supplier(std::vector<std::string> columns = {});
Review comment:
Can you add a terse description for these methods, so that non-TPC-H experts don't have to guess too much? (at least a one-liner)
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<float>(scale_factor));
+
+ ExecNode* lineitem =
+ *gen.Lineitem({"L_QUANTITY", "L_EXTENDEDPRICE", "L_TAX", "L_DISCOUNT", "L_SHIPDATE",
+ "L_RETURNFLAG", "L_LINESTATUS"});
+
+ std::shared_ptr<Date32Scalar> sept_2_1998 = std::make_shared<Date32Scalar>(
+ 10471); // September 2, 1998 is 10471 days after January 1, 1970
Review comment:
Nit, but can use `auto` to avoid repeating a long type and `const` to stress that this won't be mutated:
```suggestion
const auto sept_2_1998 = std::make_shared<Date32Scalar>(
10471); // September 2, 1998 is 10471 days after January 1, 1970
```
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
+class OrdersAndLineItemGenerator;
+class PartAndPartSupplierGenerator;
+
+class ARROW_EXPORT TpchGen {
+ public:
+ /*
+ * \brief Create a factory for nodes that generate TPC-H data
+ *
+ * Note: Individual tables will reference each other. It is important that you only
+ * create a single TpchGen instance for each plan and then you can create nodes for each
+ * table from that single TpchGen instance. Note: Every batch will be scheduled as a new
+ * task using the ExecPlan's scheduler.
+ */
+ static Result<TpchGen> Make(ExecPlan* plan, float scale_factor = 1.0f,
Review comment:
Nit: there doesn't seem to be a reason to pass a `float` here, just use a `double`?
##########
File path: cpp/src/arrow/vendored/CMakeLists.txt
##########
@@ -19,3 +19,4 @@ arrow_install_all_headers("arrow/vendored")
add_subdirectory(datetime)
add_subdirectory(double-conversion)
+add_subdirectory(pcg)
Review comment:
I _think_ we should be able to revert these changes if you're not including the PCG headers from a `.h`.
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
Review comment:
Please try to avoid including `api.h` files as they're generally costly.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
Review comment:
Put all this in the `arrow::compute::internal` namespace?
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
Review comment:
Mutable argument should be passed by pointer, not reference. Also usually they will be the last argument by position:
```suggestion
void VerifyUniqueKey(const Datum& d, int32_t min, int32_t max,
std::unordered_set<int32_t>* seen) {
```
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
Review comment:
Hmm, generally you should call the `Validate` or `ValidateFull` method on the object, rather than directly call functions from `validate.h`.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
Review comment:
Also, I'm curious: is this deliberately non-deterministic between runs? @westonpace
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row + num_offset;
+ int64_t num = 0;
+ int ibyte = static_cast<int>(num_offset);
+ for (; *num_str && ibyte < byte_width; ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
Review comment:
What is this doing? Please at least add a comment explaining. But it also seems this might be made higher level instead of parsing by hand?
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row + num_offset;
+ int64_t num = 0;
+ int ibyte = static_cast<int>(num_offset);
+ for (; *num_str && ibyte < byte_width; ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ if (verify_padding) {
+ int num_chars = ibyte - num_offset;
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ VerifyStringAndNumber_Single(row, prefix, i, nums, byte_width, verify_padding);
+ }
+}
+
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ char tmp_str[256] = {};
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ std::memcpy(tmp_str, str + start, str_len);
+ VerifyStringAndNumber_Single(tmp_str, prefix, i, nums, sizeof(tmp_str),
+ verify_padding);
+ }
+}
+
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
Review comment:
Same comment here and for all other Verify functions about visiting using higher-level constructs such as `VisitArrayDataInline`.
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row + num_offset;
+ int64_t num = 0;
+ int ibyte = static_cast<int>(num_offset);
+ for (; *num_str && ibyte < byte_width; ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ if (verify_padding) {
+ int num_chars = ibyte - num_offset;
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ VerifyStringAndNumber_Single(row, prefix, i, nums, byte_width, verify_padding);
+ }
+}
+
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ char tmp_str[256] = {};
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ std::memcpy(tmp_str, str + start, str_len);
+ VerifyStringAndNumber_Single(tmp_str, prefix, i, nums, sizeof(tmp_str),
+ verify_padding);
+ }
Review comment:
You can visit a strings array using e.g. `VisitArrayDataInline`. No need to inspect all buffers by hand. Also it seems you're not checking for null values?
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row + num_offset;
+ int64_t num = 0;
+ int ibyte = static_cast<int>(num_offset);
+ for (; *num_str && ibyte < byte_width; ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ if (verify_padding) {
+ int num_chars = ibyte - num_offset;
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ VerifyStringAndNumber_Single(row, prefix, i, nums, byte_width, verify_padding);
+ }
+}
+
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ char tmp_str[256] = {};
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ std::memcpy(tmp_str, str + start, str_len);
+ VerifyStringAndNumber_Single(tmp_str, prefix, i, nums, sizeof(tmp_str),
+ verify_padding);
+ }
+}
+
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = off[i];
+ int32_t end = off[i + 1];
+ int32_t str_len = end - start;
+ ASSERT_LE(str_len, max_length);
+ ASSERT_GE(str_len, min_length);
+ for (int32_t i = start; i < end; i++) {
+ bool is_valid =
+ std::isdigit(str[i]) || std::isalpha(str[i]) || str[i] == ',' || str[i] == ' ';
+ ASSERT_TRUE(is_valid) << "Character " << str[i]
+ << " is not a digit, a letter, a comma, or a space";
+ }
+ }
+}
+
+void VerifyModuloBetween(const Datum& d, int32_t min, int32_t max, int32_t mod) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t m = n[i] % mod;
+ ASSERT_GE(m, min) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ ASSERT_LE(m, max) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ }
+}
+
+void VerifyAllBetween(const Datum& d, int32_t min, int32_t max) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ ASSERT_GE(n[i], min) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ ASSERT_LE(n[i], max) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ }
+}
+
+void VerifyNationKey(const Datum& d) { VerifyAllBetween(d, 0, 24); }
+
+void VerifyPhone(const Datum& d) {
+ int64_t length = d.length();
+ const char* phones = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ constexpr int kByteWidth = 15; // This is common for all PHONE columns
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = phones + i * kByteWidth;
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ }
+}
+
+void VerifyDecimalsBetween(const Datum& d, int64_t min, int64_t max) {
+ int64_t length = d.length();
+ const Decimal128* decs =
+ reinterpret_cast<const Decimal128*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int64_t val = static_cast<int64_t>(decs[i]);
+ ASSERT_LE(val, max);
+ ASSERT_GE(val, min);
+ }
+}
+
+void VerifyCorrectNumberOfWords_Varlen(const Datum& d, int num_words) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+
+ int32_t start = offsets[i];
+ int32_t end = offsets[i + 1];
+ int32_t str_len = end - start;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, str + start, str_len);
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = offsets[i]; j < offsets[i + 1]; j++) {
+ bool is_space = str[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(str[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << tmp_str;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << tmp_str;
+ }
+}
+
+void VerifyCorrectNumberOfWords_FixedWidth(const Datum& d, int num_words,
+ int byte_width) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+ const char* row = str + i * byte_width;
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = 0; j < byte_width && row[j]; j++) {
+ bool is_space = row[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(row[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << row;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << row;
+ }
+}
+
+void VerifyOneOf(const Datum& d, const std::unordered_set<char>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++)
+ ASSERT_TRUE(possibilities.find(col[i]) != possibilities.end());
+}
+
+void VerifyOneOf(const Datum& d, int32_t byte_width,
+ const std::unordered_set<std::string>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = col + i * byte_width;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, row, byte_width);
+ ASSERT_TRUE(possibilities.find(tmp_str) != possibilities.end())
+ << tmp_str << " is not a valid string.";
+ }
+}
+
+void CountInstances(std::unordered_map<int32_t, int32_t>& counts, const Datum& d) {
+ int64_t length = d.length();
+ const int32_t* nums = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) counts[nums[i]]++;
+}
+
+void CountModifiedComments(const Datum& d, int& good_count, int& bad_count) {
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ // Length of S_COMMENT is at most 100
+ char tmp_string[101];
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + offsets[i];
+ int32_t row_length = offsets[i + 1] - offsets[i];
+ std::memset(tmp_string, 0, sizeof(tmp_string));
+ std::memcpy(tmp_string, row, row_length);
+ char* customer = std::strstr(tmp_string, "Customer");
+ char* recommends = std::strstr(tmp_string, "Recommends");
+ char* complaints = std::strstr(tmp_string, "Complaints");
+ if (customer) {
+ ASSERT_TRUE((recommends != nullptr) ^ (complaints != nullptr));
+ if (recommends) good_count++;
+ if (complaints) bad_count++;
+ }
+ }
+}
+
+TEST(TpchNode, ScaleFactor) {
+ ExecContext ctx(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(&ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), 0.25f);
+ ExecNode* table = *gen.Supplier();
+ AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+ Declaration sink("sink", {Declaration::Input(table)}, SinkNodeOptions{&sink_gen});
+ std::ignore = *sink.AddToPlan(plan.get());
+ auto fut = StartAndCollect(plan.get(), sink_gen);
+ auto res = *fut.MoveResult();
+
+ int64_t kExpectedRows = 2500;
+ int64_t num_rows = 0;
+ for (auto& batch : res) num_rows += batch.length;
+ ASSERT_EQ(num_rows, kExpectedRows);
+ arrow::internal::GetCpuThreadPool()->WaitForIdle();
Review comment:
Why is this necessary? Can you add a comment here?
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row + num_offset;
+ int64_t num = 0;
+ int ibyte = static_cast<int>(num_offset);
+ for (; *num_str && ibyte < byte_width; ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ if (verify_padding) {
+ int num_chars = ibyte - num_offset;
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ VerifyStringAndNumber_Single(row, prefix, i, nums, byte_width, verify_padding);
+ }
+}
+
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ char tmp_str[256] = {};
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ std::memcpy(tmp_str, str + start, str_len);
+ VerifyStringAndNumber_Single(tmp_str, prefix, i, nums, sizeof(tmp_str),
+ verify_padding);
+ }
+}
+
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = off[i];
+ int32_t end = off[i + 1];
+ int32_t str_len = end - start;
+ ASSERT_LE(str_len, max_length);
+ ASSERT_GE(str_len, min_length);
+ for (int32_t i = start; i < end; i++) {
+ bool is_valid =
+ std::isdigit(str[i]) || std::isalpha(str[i]) || str[i] == ',' || str[i] == ' ';
+ ASSERT_TRUE(is_valid) << "Character " << str[i]
+ << " is not a digit, a letter, a comma, or a space";
+ }
+ }
+}
+
+void VerifyModuloBetween(const Datum& d, int32_t min, int32_t max, int32_t mod) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t m = n[i] % mod;
+ ASSERT_GE(m, min) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ ASSERT_LE(m, max) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ }
+}
+
+void VerifyAllBetween(const Datum& d, int32_t min, int32_t max) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ ASSERT_GE(n[i], min) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ ASSERT_LE(n[i], max) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ }
+}
+
+void VerifyNationKey(const Datum& d) { VerifyAllBetween(d, 0, 24); }
+
+void VerifyPhone(const Datum& d) {
+ int64_t length = d.length();
+ const char* phones = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ constexpr int kByteWidth = 15; // This is common for all PHONE columns
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = phones + i * kByteWidth;
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ }
+}
+
+void VerifyDecimalsBetween(const Datum& d, int64_t min, int64_t max) {
+ int64_t length = d.length();
+ const Decimal128* decs =
+ reinterpret_cast<const Decimal128*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int64_t val = static_cast<int64_t>(decs[i]);
+ ASSERT_LE(val, max);
+ ASSERT_GE(val, min);
+ }
+}
+
+void VerifyCorrectNumberOfWords_Varlen(const Datum& d, int num_words) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+
+ int32_t start = offsets[i];
+ int32_t end = offsets[i + 1];
+ int32_t str_len = end - start;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, str + start, str_len);
Review comment:
Please let's refrain from using C-style code, you can use e.g. `util::string_view` to handle a range of bytes as a string.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
Review comment:
Hmm, unless there's a strong reason to have a module global like this, it is generally better to have a local variable or member attribute where needed.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
Review comment:
Can you put all private declarations in the anonymous space? That will hide them from outside code and then you don't have to add `static` to some declarations.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ if (!partsupp_output_queue_.empty()) {
+ ExecBatch result = std::move(partsupp_output_queue_.front());
+ partsupp_output_queue_.pop();
+ return std::move(result);
+ }
+ }
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+ if (!part_cols_.empty()) {
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch part_batch,
+ ExecBatch::Make(std::move(part_result)));
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ part_output_queue_.emplace(std::move(part_batch));
+ }
+ }
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ for (size_t i = 1; i < partsupp_results.size(); i++)
+ partsupp_output_queue_.emplace(std::move(partsupp_results[i]));
+ }
+ return std::move(partsupp_results[0]);
+ }
+
+ private:
+#define FOR_EACH_PART_COLUMN(F) \
+ F(P_PARTKEY) \
+ F(P_NAME) \
+ F(P_MFGR) \
+ F(P_BRAND) \
+ F(P_TYPE) \
+ F(P_SIZE) \
+ F(P_CONTAINER) \
+ F(P_RETAILPRICE) \
+ F(P_COMMENT)
+
+#define FOR_EACH_PARTSUPP_COLUMN(F) \
+ F(PS_PARTKEY) \
+ F(PS_SUPPKEY) \
+ F(PS_AVAILQTY) \
+ F(PS_SUPPLYCOST) \
+ F(PS_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct PART {
+ enum {
+ FOR_EACH_PART_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct PARTSUPP {
+ enum {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, PART::col},
+ const std::unordered_map<std::string, int> kPartNameMap = {
+ FOR_EACH_PART_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, PARTSUPP::col},
+ const std::unordered_map<std::string, int> kPartsuppNameMap = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kPartGenerators = {FOR_EACH_PART_COLUMN(MAKE_FN_ARRAY)};
+ std::vector<GenerateColumnFn> kPartsuppGenerators = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kPartTypes = {
+ int32(), utf8(), fixed_size_binary(25), fixed_size_binary(10),
+ utf8(), int32(), fixed_size_binary(10), decimal(12, 2),
+ utf8(),
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kPartsuppTypes = {
+ int32(), int32(), int32(), decimal(12, 2), utf8(),
+ };
+
+ Status AllocatePartBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.part[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.part_to_generate * byte_width));
+ ArrayData ad(kPartTypes[column], tld.part_to_generate, {nullptr, std::move(buff)});
+ tld.part[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status P_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_PARTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_PARTKEY));
+ int32_t* p_partkey = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) {
+ p_partkey[i] = static_cast<int32_t>(tld.partkey_start + i + 1);
+ ARROW_DCHECK(1 <= p_partkey[i] && p_partkey[i] <= part_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_NAME].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(
+ 0, static_cast<uint8_t>(kNumNameParts - 1));
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dist(tld.rng));
+ tld.string_indices[irow * 5 + ipart] = name_part_index;
+ string_length += std::strlen(NameParts[name_part_index]);
+ }
+ // Add 4 because there is a space between each word (i.e. four spaces)
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length + 4);
+ }
+ // Add an extra byte for the space after in the very last string.
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate] + 1));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 5 + ipart];
+ const char* part = NameParts[name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ *row++ = ' ';
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_NAME], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_NAME] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_MFGR(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_MFGR].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* manufacturer = "Manufacturer#";
+ const size_t manufacturer_length = std::strlen(manufacturer);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_MFGR));
+ char* p_mfgr = reinterpret_cast<char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ std::strncpy(p_mfgr + byte_width * irow, manufacturer, byte_width);
+ char mfgr_id = '0' + dist(tld.rng);
+ *(p_mfgr + byte_width * irow + manufacturer_length) = mfgr_id;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_BRAND(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_BRAND].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_MFGR(thread_index));
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* brand = "Brand#";
+ const size_t brand_length = std::strlen(brand);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_BRAND));
+ const char* p_mfgr = reinterpret_cast<const char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->data());
+ char* p_brand = reinterpret_cast<char*>(
+ tld.part[PART::P_BRAND].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_BRAND]);
+ int32_t mfgr_byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ const size_t mfgr_id_offset = std::strlen("Manufacturer#");
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = p_brand + byte_width * irow;
+ char mfgr_id = *(p_mfgr + irow * mfgr_byte_width + mfgr_id_offset);
+ char brand_id = '0' + dist(tld.rng);
+ std::strncpy(row, brand, byte_width);
+ *(row + brand_length) = mfgr_id;
+ *(row + brand_length + 1) = brand_id;
+ irow += 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_TYPE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_TYPE].kind() == Datum::NONE) {
+ using D = std::uniform_int_distribution<uint32_t>;
+ D dists[] = {
+ D{0, static_cast<uint8_t>(kNumTypes_1 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_2 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_3 - 1)},
+ };
+
+ const char** types[] = {Types_1, Types_2, Types_3};
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dists[ipart](tld.rng));
+ tld.string_indices[irow * 3 + ipart] = name_part_index;
+ string_length += std::strlen(types[ipart][name_part_index]);
+ }
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate]));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 3 + ipart];
+ const char* part = types[ipart][name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_TYPE], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_TYPE] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_SIZE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_SIZE].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(1, 50);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_SIZE));
+ int32_t* p_size = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_SIZE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) p_size[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status P_CONTAINER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_CONTAINER].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist1(
+ 0, static_cast<uint8_t>(kNumContainers_1 - 1));
+ std::uniform_int_distribution<int> dist2(
+ 0, static_cast<uint8_t>(kNumContainers_2 - 1));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_CONTAINER));
+ char* p_container = reinterpret_cast<char*>(
+ tld.part[PART::P_CONTAINER].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_CONTAINER]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int container1_idx = dist1(tld.rng);
+ int container2_idx = dist2(tld.rng);
+ const char* container1 = Containers_1[container1_idx];
+ const char* container2 = Containers_2[container2_idx];
+ size_t container1_length = std::strlen(container1);
+ size_t container2_length = std::strlen(container2);
+
+ char* row = p_container + byte_width * irow;
+ std::strncpy(row, container1, byte_width);
+ std::memcpy(row + container1_length, container2, container2_length);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_RETAILPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_RETAILPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_RETAILPRICE));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+ Decimal128* p_retailprice = reinterpret_cast<Decimal128*>(
+ tld.part[PART::P_RETAILPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int32_t partkey = p_partkey[irow];
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ p_retailprice[irow] = {retail_price};
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.part[PART::P_COMMENT],
+ g_text.GenerateComments(tld.part_to_generate, 5, 22, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ int64_t PartsuppBatchesToGenerate(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t num_batches = (ps_to_generate + batch_size_ - 1) / batch_size_;
+ return num_batches;
+ }
+
+ Status InitPartsupp(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.generated_partsupp.reset();
+ int64_t num_batches = PartsuppBatchesToGenerate(thread_index);
+ tld.partsupp.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.partsupp) {
+ batch.resize(PARTSUPP::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocatePartSuppBatch(size_t thread_index, size_t ibatch, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartsuppTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kPartsuppTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.partsupp[ibatch][column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status PS_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_PARTKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_PARTKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_PARTKEY));
+ int32_t* ps_partkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run; ipartsupp++, irun++)
+ ps_partkey[batch_offset++] = p_partkey[ipart];
+
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPKEY));
+ int32_t* ps_suppkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run;
+ ipartsupp++, irun++) {
+ int32_t supplier = static_cast<int32_t>(ipartsupp);
+ int32_t partkey = p_partkey[ipart];
+ ps_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_AVAILQTY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_AVAILQTY]) {
+ tld.generated_partsupp[PARTSUPP::PS_AVAILQTY] = true;
+ std::uniform_int_distribution<int32_t> dist(1, 9999);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_AVAILQTY));
+ int32_t* ps_availqty =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for (int64_t irun = 0; irun < next_run; irun++) ps_availqty[irun] = dist(tld.rng);
+
+ tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPLYCOST(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST] = true;
+ std::uniform_int_distribution<int64_t> dist(100, 100000);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(
+ AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPLYCOST));
+ Decimal128* ps_supplycost =
+ reinterpret_cast<Decimal128*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for (int64_t irun = 0; irun < next_run; irun++)
+ ps_supplycost[irun] = {dist(tld.rng)};
+
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PARTSUPP::PS_COMMENT].kind() == Datum::NONE) {
+ int64_t irow = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ int64_t num_rows = std::min(batch_size_, ps_to_generate - irow);
+ ARROW_ASSIGN_OR_RAISE(tld.partsupp[ibatch][PARTSUPP::PS_COMMENT],
+ g_text.GenerateComments(num_rows, 49, 198, tld.rng));
+ irow += num_rows;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ std::vector<Datum> part;
+ std::vector<int8_t> string_indices;
+ int64_t part_to_generate{0};
+ int64_t partkey_start{0};
+
+ std::vector<std::vector<Datum>> partsupp;
+ std::bitset<PARTSUPP::kNumCols> generated_partsupp;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex part_output_queue_mutex_;
+ std::mutex partsupp_output_queue_mutex_;
+ std::queue<ExecBatch> part_output_queue_;
+ std::queue<ExecBatch> partsupp_output_queue_;
+ int64_t batch_size_{0};
+ float scale_factor_{0};
+ int64_t part_rows_to_generate_{0};
+ int64_t part_rows_generated_{0};
+ std::vector<int> part_cols_;
+ std::vector<int> partsupp_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<int64_t> part_batches_generated_ = {0};
+ std::atomic<int64_t> partsupp_batches_generated_ = {0};
+ static constexpr int64_t kPartSuppRowsPerPart = 4;
+};
+
+class OrdersAndLineItemGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ tld.items_per_order.resize(batch_size_);
+ tld.rng.seed(seq);
+ }
+ orders_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000 * 10);
+ }
+ return Status::OK();
+ }
+
+ int64_t orders_batches_generated() const { return orders_batches_generated_.load(); }
+
+ int64_t lineitem_batches_generated() const {
+ return lineitem_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetOrdersOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kOrdersTypes, kOrdersNameMap, orders_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetLineItemOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kLineitemTypes, kLineitemNameMap, lineitem_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextOrdersBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if (!orders_output_queue_.empty()) {
+ ExecBatch batch = std::move(orders_output_queue_.front());
+ orders_output_queue_.pop();
+ return std::move(batch);
+ } else if (orders_rows_generated_ == orders_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate =
+ std::min(batch_size_, orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.first_batch_offset = 0;
+ tld.generated_lineitem.reset();
+
+ for (int col : orders_cols_) RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for (int col : lineitem_cols_) RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for (size_t i = 0; i < orders_cols_.size(); i++) {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ if (!lineitem_cols_.empty()) {
+ std::vector<ExecBatch> lineitem_results;
+ for (size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++) {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for (size_t icol = 0; icol < lineitem_cols_.size(); icol++) {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(lineitem_output_queue_mutex_);
+ for (ExecBatch& eb : lineitem_results) {
+ lineitem_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(orders_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextLineItemBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ExecBatch queued;
+ bool from_queue = false;
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ if (!lineitem_output_queue_.empty()) {
+ queued = std::move(lineitem_output_queue_.front());
+ lineitem_output_queue_.pop();
+ from_queue = true;
+ }
+ }
+ tld.first_batch_offset = 0;
+ if (from_queue) {
+ ARROW_DCHECK(queued.length <= batch_size_);
+ tld.first_batch_offset = queued.length;
+ if (queued.length == batch_size_) return std::move(queued);
+ }
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if (orders_rows_generated_ == orders_rows_to_generate_) {
+ if (from_queue) return std::move(queued);
+ return util::nullopt;
+ }
+
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate =
+ std::min(batch_size_, orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1ll);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.generated_lineitem.reset();
+ if (from_queue) {
+ lineitem_batches_generated_.fetch_sub(1);
+ for (size_t i = 0; i < lineitem_cols_.size(); i++)
+ if (tld.lineitem[0][lineitem_cols_[i]].kind() == Datum::NONE)
+ tld.lineitem[0][lineitem_cols_[i]] = std::move(queued[i]);
+ }
+
+ for (int col : orders_cols_) RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for (int col : lineitem_cols_) RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ if (!orders_cols_.empty()) {
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for (size_t i = 0; i < orders_cols_.size(); i++) {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch orders_batch,
+ ExecBatch::Make(std::move(orders_result)));
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ orders_output_queue_.emplace(std::move(orders_batch));
+ }
+ }
+ std::vector<ExecBatch> lineitem_results;
+ for (size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++) {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for (size_t icol = 0; icol < lineitem_cols_.size(); icol++) {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ lineitem_batches_generated_.fetch_add(static_cast<int64_t>(lineitem_results.size()));
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ for (size_t i = 1; i < lineitem_results.size(); i++)
+ lineitem_output_queue_.emplace(std::move(lineitem_results[i]));
+ }
+ return std::move(lineitem_results[0]);
+ }
+
+ private:
+#define FOR_EACH_ORDERS_COLUMN(F) \
+ F(O_ORDERKEY) \
+ F(O_CUSTKEY) \
+ F(O_ORDERSTATUS) \
+ F(O_TOTALPRICE) \
+ F(O_ORDERDATE) \
+ F(O_ORDERPRIORITY) \
+ F(O_CLERK) \
+ F(O_SHIPPRIORITY) \
+ F(O_COMMENT)
+
+#define FOR_EACH_LINEITEM_COLUMN(F) \
+ F(L_ORDERKEY) \
+ F(L_PARTKEY) \
+ F(L_SUPPKEY) \
+ F(L_LINENUMBER) \
+ F(L_QUANTITY) \
+ F(L_EXTENDEDPRICE) \
+ F(L_DISCOUNT) \
+ F(L_TAX) \
+ F(L_RETURNFLAG) \
+ F(L_LINESTATUS) \
+ F(L_SHIPDATE) \
+ F(L_COMMITDATE) \
+ F(L_RECEIPTDATE) \
+ F(L_SHIPINSTRUCT) \
+ F(L_SHIPMODE) \
+ F(L_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct ORDERS {
+ enum {
+ FOR_EACH_ORDERS_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct LINEITEM {
+ enum {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, ORDERS::col},
+ const std::unordered_map<std::string, int> kOrdersNameMap = {
+ FOR_EACH_ORDERS_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, LINEITEM::col},
+ const std::unordered_map<std::string, int> kLineitemNameMap = {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ const std::vector<GenerateColumnFn> kOrdersGenerators = {
+ FOR_EACH_ORDERS_COLUMN(MAKE_FN_ARRAY)};
+ const std::vector<GenerateColumnFn> kLineitemGenerators = {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kOrdersTypes = {int32(),
+ int32(),
+ fixed_size_binary(1),
+ decimal(12, 2),
+ date32(),
+ fixed_size_binary(15),
+ fixed_size_binary(15),
+ int32(),
+ utf8()};
+
+ const std::vector<std::shared_ptr<DataType>> kLineitemTypes = {
+ int32(),
+ int32(),
+ int32(),
+ int32(),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ fixed_size_binary(1),
+ fixed_size_binary(1),
+ date32(),
+ date32(),
+ date32(),
+ fixed_size_binary(25),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status AllocateOrdersBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.orders[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.orders_to_generate * byte_width));
+ ArrayData ad(kOrdersTypes[column], tld.orders_to_generate,
+ {nullptr, std::move(buff)});
+ tld.orders[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status O_ORDERKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERKEY));
+ int32_t* o_orderkey = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ int32_t orderkey_index = static_cast<int32_t>(tld.orderkey_start + i);
+ int32_t index_of_run = orderkey_index / 8;
+ int32_t index_in_run = orderkey_index % 8;
+ o_orderkey[i] = (index_of_run * 32 + index_in_run + 1);
+ ARROW_DCHECK(1 <= o_orderkey[i] && o_orderkey[i] <= 4 * orders_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CUSTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_CUSTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CUSTKEY));
+
+ // Spec says it must be a random number between 1 and SF*150000 that is not
+ // divisible by 3. Rather than repeatedly generating numbers until we get to
+ // a non-divisible-by-3 number, we just generate a number between
+ // 0 and SF * 50000 - 1, multiply by 3, and then add either 1 or 2.
+ int32_t sf_50k = static_cast<int32_t>(scale_factor_ * 50000);
+ std::uniform_int_distribution<int32_t> base_dist(0, sf_50k - 1);
+ std::uniform_int_distribution<int32_t> offset_dist(1, 2);
+ int32_t* o_custkey = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_CUSTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++)
+ o_custkey[i] = 3 * base_dist(tld.rng) + offset_dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERSTATUS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERSTATUS].kind() == Datum::NONE) {
+ RETURN_NOT_OK(L_LINESTATUS(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERSTATUS));
+
+ char* o_orderstatus = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_ORDERSTATUS].array()->buffers[1]->mutable_data());
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ bool all_f = true;
+ bool all_o = true;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ const char* l_linestatus = reinterpret_cast<const char*>(
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->buffers[1]->data());
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++, batch_offset++) {
+ all_f &= l_linestatus[batch_offset] == 'F';
+ all_o &= l_linestatus[batch_offset] == 'O';
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ ARROW_DCHECK(!(all_f && all_o));
+ if (all_f)
+ o_orderstatus[iorder] = 'F';
+ else if (all_o)
+ o_orderstatus[iorder] = 'O';
+ else
+ o_orderstatus[iorder] = 'P';
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_TOTALPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_TOTALPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(L_EXTENDEDPRICE(thread_index));
+ RETURN_NOT_OK(L_TAX(thread_index));
+ RETURN_NOT_OK(L_DISCOUNT(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_TOTALPRICE));
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ int64_t sum = 0;
+ Decimal128* o_totalprice = reinterpret_cast<Decimal128*>(
+ tld.orders[ORDERS::O_TOTALPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ const Decimal128* l_extendedprice = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->buffers[1]->data());
+ const Decimal128* l_tax = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->data());
+ const Decimal128* l_discount = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->buffers[1]->data());
+
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++, batch_offset++) {
+ int64_t eprice = static_cast<int64_t>(l_extendedprice[batch_offset]);
+ int64_t tax = static_cast<int64_t>(l_tax[batch_offset]);
+ int64_t discount = static_cast<int64_t>(l_discount[batch_offset]);
+ sum += (eprice * (100 + tax) * (100 - discount));
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ sum /= 100 * 100;
+ o_totalprice[iorder] = {sum};
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERDATE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERDATE));
+
+ std::uniform_int_distribution<uint32_t> dist(kStartDate, kEndDate - 151);
+ uint32_t* o_orderdate = reinterpret_cast<uint32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) o_orderdate[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERPRIORITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERPRIORITY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERPRIORITY));
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_ORDERPRIORITY]);
+ std::uniform_int_distribution<int32_t> dist(0, kNumPriorities - 1);
+ char* o_orderpriority = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_ORDERPRIORITY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ const char* str = Priorities[dist(tld.rng)];
+ std::strncpy(o_orderpriority + i * byte_width, str, byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CLERK(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_CLERK].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CLERK));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_CLERK]);
+ int64_t max_clerk_id = static_cast<int64_t>(scale_factor_ * 1000);
+ std::uniform_int_distribution<int64_t> dist(1, max_clerk_id);
+ char* o_clerk = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_CLERK].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ const char* clerk = "Clerk#";
+ const size_t clerk_length = std::strlen(clerk);
+ int64_t clerk_number = dist(tld.rng);
+ char* output = o_clerk + i * byte_width;
+ std::strncpy(output, clerk, byte_width);
+ AppendNumberPaddedToNineDigits(output + clerk_length, clerk_number);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_SHIPPRIORITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_SHIPPRIORITY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_SHIPPRIORITY));
+ int32_t* o_shippriority = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_SHIPPRIORITY].array()->buffers[1]->mutable_data());
+ std::memset(o_shippriority, 0, tld.orders_to_generate * sizeof(int32_t));
+ }
+ return Status::OK();
+ }
+
+ Status O_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.orders[ORDERS::O_COMMENT],
+ g_text.GenerateComments(tld.orders_to_generate, 19, 78, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ Status GenerateRowCounts(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ std::uniform_int_distribution<int> length_dist(1, 7);
+ tld.lineitem_to_generate = 0;
+ tld.items_per_order.clear();
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ int length = length_dist(tld.rng);
+ tld.items_per_order.push_back(length);
+ tld.lineitem_to_generate += length;
+ }
+ int64_t num_batches =
+ (tld.first_batch_offset + tld.lineitem_to_generate + batch_size_ - 1) /
+ batch_size_;
+ tld.lineitem.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.lineitem) {
+ batch.resize(LINEITEM::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocateLineItemBufferIfNeeded(size_t thread_index, size_t ibatch, int column,
+ size_t& out_batch_offset) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.lineitem[ibatch][column].kind() == Datum::NONE) {
+ int32_t byte_width = arrow::internal::GetByteWidth(*kLineitemTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kLineitemTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.lineitem[ibatch][column] = std::move(ad);
+ out_batch_offset = 0;
+ }
+ if (ibatch == 0) out_batch_offset = tld.first_batch_offset;
+
+ return Status::OK();
+ }
+
+ Status L_ORDERKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_ORDERKEY]) {
+ tld.generated_lineitem[LINEITEM::L_ORDERKEY] = true;
+ RETURN_NOT_OK(O_ORDERKEY(thread_index));
+ const int32_t* o_orderkey = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_ORDERKEY, batch_offset));
+ int32_t* l_linenumber =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_ORDERKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_linenumber[batch_offset++] = o_orderkey[iorder];
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_ORDERKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_PARTKEY]) {
+ tld.generated_lineitem[LINEITEM::L_PARTKEY] = true;
+
+ size_t ibatch = 0;
+ int32_t max_partkey = static_cast<int32_t>(scale_factor_ * 200000);
+ std::uniform_int_distribution<int32_t> dist(1, max_partkey);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_PARTKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_partkey =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_partkey[batch_offset] = dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SUPPKEY]) {
+ tld.generated_lineitem[LINEITEM::L_SUPPKEY] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 3);
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset = 0;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SUPPKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_suppkey =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const int32_t* l_partkey = reinterpret_cast<const int32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ for (int64_t i = 0; i < next_run; i++) {
+ int32_t supplier = dist(tld.rng);
+ int32_t partkey = l_partkey[batch_offset];
+ // Fun fact: the parentheses for this expression are unbalanced in the TPC-H
+ // spec.
+ l_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SUPPKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINENUMBER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_LINENUMBER]) {
+ tld.generated_lineitem[LINEITEM::L_LINENUMBER] = true;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_LINENUMBER, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_linenumber =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_LINENUMBER]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++) {
+ l_linenumber[batch_offset++] = (iline + 1);
+ ARROW_DCHECK(1 <= (iline + 1) && (iline + 1) <= 7);
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINENUMBER].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_QUANTITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_QUANTITY]) {
+ tld.generated_lineitem[LINEITEM::L_QUANTITY] = true;
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int64_t> dist(1, 50);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_QUANTITY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_quantity =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_QUANTITY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t i = 0; i < next_run; i++) {
+ // Multiply by 100 because the type is decimal(12, 2), so the decimal goes after
+ // two digits
+ int64_t quantity = dist(tld.rng) * 100;
+ l_quantity[batch_offset++] = {quantity};
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_EXTENDEDPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE]) {
+ tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+ RETURN_NOT_OK(L_QUANTITY(thread_index));
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_EXTENDEDPRICE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ const int32_t* l_partkey = reinterpret_cast<const int32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ const Decimal128* l_quantity = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->buffers[1]->data());
+ Decimal128* l_extendedprice =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ int64_t partkey = static_cast<int64_t>(l_partkey[batch_offset]);
+ // Divide by 100 to recover the integer representation (not Decimal).
+ int64_t quantity = static_cast<int64_t>(l_quantity[batch_offset]) / 100;
+
+ // Spec says to divide by 100, but that happens automatically due to this being
+ // stored to two decimal points.
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ int64_t extended_price = retail_price * quantity;
+ l_extendedprice[batch_offset] = {extended_price};
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_DISCOUNT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_DISCOUNT]) {
+ tld.generated_lineitem[LINEITEM::L_DISCOUNT] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 10);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_DISCOUNT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_discount =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_DISCOUNT]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_discount[batch_offset] = {dist(tld.rng)};
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_TAX(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_TAX]) {
+ tld.generated_lineitem[LINEITEM::L_TAX] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 8);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_TAX, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_tax = reinterpret_cast<Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_tax[batch_offset] = {dist(tld.rng)};
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RETURNFLAG(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_RETURNFLAG]) {
+ tld.generated_lineitem[LINEITEM::L_RETURNFLAG] = true;
+ RETURN_NOT_OK(L_RECEIPTDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<uint32_t> dist;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_RETURNFLAG, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_returnflag =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_receiptdate = reinterpret_cast<const uint32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ if (l_receiptdate[batch_offset] <= kCurrentDate) {
+ uint32_t r = dist(tld.rng);
+ l_returnflag[batch_offset] = (r % 2 == 1) ? 'R' : 'A';
+ } else {
+ l_returnflag[batch_offset] = 'N';
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINESTATUS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_LINESTATUS]) {
+ tld.generated_lineitem[LINEITEM::L_LINESTATUS] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_LINESTATUS, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_linestatus =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_LINESTATUS]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_shipdate =
+ reinterpret_cast<const uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ if (l_shipdate[batch_offset] > kCurrentDate)
+ l_linestatus[batch_offset] = 'O';
+ else
+ l_linestatus[batch_offset] = 'F';
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPDATE]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPDATE] = true;
+ RETURN_NOT_OK(O_ORDERDATE(thread_index));
+ const int32_t* o_orderdate = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(1, 121);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SHIPDATE, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_shipdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_shipdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMITDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_COMMITDATE]) {
+ tld.generated_lineitem[LINEITEM::L_COMMITDATE] = true;
+ const int32_t* o_orderdate = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(30, 90);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_COMMITDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_commitdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_COMMITDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_commitdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMITDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RECEIPTDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_RECEIPTDATE]) {
+ tld.generated_lineitem[LINEITEM::L_RECEIPTDATE] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(1, 30);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_RECEIPTDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_receiptdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_shipdate =
+ reinterpret_cast<const uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_receiptdate[batch_offset] = l_shipdate[batch_offset] + dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPINSTRUCT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT] = true;
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPINSTRUCT]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumInstructions - 1);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_SHIPINSTRUCT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_shipinstruct =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ const char* str = Instructions[dist(tld.rng)];
+ // Note that we don't have to memset the buffer to 0 because strncpy pads each
+ // string with 0's anyway
+ std::strncpy(l_shipinstruct + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPMODE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPMODE]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPMODE] = true;
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPMODE]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumModes - 1);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SHIPMODE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_shipmode =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_SHIPMODE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ const char* str = Modes[dist(tld.rng)];
+ std::strncpy(l_shipmode + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPMODE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_COMMENT]) {
+ tld.generated_lineitem[LINEITEM::L_COMMENT] = true;
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ // Comments are kind of sneaky: we always generate the full batch and then just
+ // bump the length
+ if (tld.lineitem[ibatch][LINEITEM::L_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.lineitem[ibatch][LINEITEM::L_COMMENT],
+ g_text.GenerateComments(batch_size_, 10, 43, tld.rng));
+ batch_offset = 0;
+ }
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ batch_offset += next_run;
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMENT].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ std::vector<Datum> orders;
+ int64_t orders_to_generate;
+ int64_t orderkey_start;
+
+ std::vector<std::vector<Datum>> lineitem;
+ std::vector<int> items_per_order;
+ int64_t lineitem_to_generate;
+ int64_t first_batch_offset;
+ std::bitset<LINEITEM::kNumCols> generated_lineitem;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex orders_output_queue_mutex_;
+ std::mutex lineitem_output_queue_mutex_;
+ std::queue<ExecBatch> orders_output_queue_;
+ std::queue<ExecBatch> lineitem_output_queue_;
+ int64_t batch_size_;
+ float scale_factor_;
+ int64_t orders_rows_to_generate_;
+ int64_t orders_rows_generated_;
+ std::vector<int> orders_cols_;
+ std::vector<int> lineitem_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<size_t> orders_batches_generated_ = {0};
+ std::atomic<size_t> lineitem_batches_generated_ = {0};
+};
+
+class SupplierGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 10000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, gen_list_));
+
+ random::pcg32_fast rng;
+ std::uniform_int_distribution<int64_t> dist(0, rows_to_generate_ - 1);
+ size_t num_special_rows = static_cast<size_t>(5 * scale_factor_);
+ std::unordered_set<int64_t> good_rows_set;
+ while (good_rows_set.size() < num_special_rows) {
+ int64_t row = dist(rng);
+ good_rows_set.insert(row);
+ }
+ std::unordered_set<int64_t> bad_rows_set;
+ while (bad_rows_set.size() < num_special_rows) {
+ int64_t bad_row;
+ do {
+ bad_row = dist(rng);
+ } while (good_rows_set.find(bad_row) != good_rows_set.end());
+ bad_rows_set.insert(bad_row);
+ }
+ good_rows_.clear();
+ bad_rows_.clear();
+ good_rows_.insert(good_rows_.end(), good_rows_set.begin(), good_rows_set.end());
+ bad_rows_.insert(bad_rows_.end(), bad_rows_set.begin(), bad_rows_set.end());
+ std::sort(good_rows_.begin(), good_rows_.end());
+ std::sort(bad_rows_.begin(), bad_rows_.end());
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) tld.rng.seed(seq);
+
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(S_SUPPKEY) \
+ F(S_NAME) \
+ F(S_ADDRESS) \
+ F(S_NATIONKEY) \
+ F(S_PHONE) \
+ F(S_ACCTBAL) \
+ F(S_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct SUPPLIER {
+ enum {
+ FOR_EACH_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) {#col, SUPPLIER::col},
+ const std::unordered_map<std::string, int> kNameMap = {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kGenerators = {FOR_EACH_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(), fixed_size_binary(25), utf8(),
+ int32(), fixed_size_binary(15), decimal(12, 2),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index) {
+ if (done_.load()) return Status::OK();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.suppkey_start = rows_generated_.fetch_add(batch_size_);
+ if (tld.suppkey_start >= rows_to_generate_) return Status::OK();
+
+ tld.to_generate = std::min(batch_size_, rows_to_generate_ - tld.suppkey_start);
+
+ tld.batch.resize(SUPPLIER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for (int col : gen_list_) RETURN_NOT_OK(kGenerators[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for (size_t i = 0; i < gen_list_.size(); i++) {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_outputted_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_outputted_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if (is_last_batch) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ return Status::OK();
+ }
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(kTypes[column], tld.to_generate, {nullptr, std::move(buff)});
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status S_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_SUPPKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_SUPPKEY));
+ int32_t* s_suppkey = reinterpret_cast<int32_t*>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ s_suppkey[irow] = static_cast<int32_t>(tld.suppkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_NAME].kind() == Datum::NONE) {
+ RETURN_NOT_OK(S_SUPPKEY(thread_index));
+ const int32_t* s_suppkey = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->data());
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NAME));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[SUPPLIER::S_NAME]);
+ char* s_name = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_NAME].array()->buffers[1]->mutable_data());
+ // Look man, I'm just following the spec ok? Section 4.2.3 as of March 1 2022
+ const char* supplier = "Supplie#r";
+ const size_t supplier_length = std::strlen(supplier);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = s_name + byte_width * irow;
+ std::strncpy(out, supplier, byte_width);
+ AppendNumberPaddedToNineDigits(out + supplier_length, s_suppkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ADDRESS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_ADDRESS].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[SUPPLIER::S_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status S_NATIONKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_NATIONKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t* s_nationkey = reinterpret_cast<int32_t*>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status S_PHONE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_PHONE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(S_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[SUPPLIER::S_PHONE]);
+ const int32_t* s_nationkey = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->data());
+ char* s_phone = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_PHONE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ GeneratePhoneNumber(s_phone + irow * byte_width, tld.rng, s_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ACCTBAL(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_ACCTBAL].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_ACCTBAL));
+ Decimal128* s_acctbal = reinterpret_cast<Decimal128*>(
+ tld.batch[SUPPLIER::S_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_acctbal[irow] = {dist(tld.rng)};
+ }
+ return Status::OK();
+ }
+
+ Status S_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[SUPPLIER::S_COMMENT],
+ g_text.GenerateComments(tld.to_generate, 25, 100, tld.rng));
+ ModifyComments(thread_index, "Recommends", good_rows_);
+ ModifyComments(thread_index, "Complaints", bad_rows_);
+ }
+ return Status::OK();
+ }
+
+ void ModifyComments(size_t thread_index, const char* review,
+ const std::vector<int64_t>& indices) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ const int32_t* offsets = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[1]->data());
+ char* str = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[2]->mutable_data());
+ const char* customer = "Customer";
+ const int32_t customer_length = static_cast<int32_t>(std::strlen(customer));
+ const int32_t review_length = static_cast<int32_t>(std::strlen(review));
+
+ auto it = std::lower_bound(indices.begin(), indices.end(), tld.suppkey_start);
+ for (; it != indices.end() && *it < tld.suppkey_start + tld.to_generate; it++) {
+ int64_t idx_in_batch = *it - tld.suppkey_start;
+ char* out = str + offsets[idx_in_batch];
+ int32_t str_length = offsets[idx_in_batch + 1] - offsets[idx_in_batch];
+ std::uniform_int_distribution<int32_t> gap_dist(
+ 0, str_length - customer_length - review_length);
+ int32_t gap = gap_dist(tld.rng);
+ int32_t total_length = customer_length + gap + review_length;
+ std::uniform_int_distribution<int32_t> start_dist(0, str_length - total_length);
+ int32_t start = start_dist(tld.rng);
+ std::memcpy(out + start, customer, customer_length);
+ std::memcpy(out + start + customer_length + gap, review, review_length);
+ }
+ }
+
+ struct ThreadLocalData {
+ random::pcg32_fast rng;
+ int64_t suppkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+ std::vector<int64_t> good_rows_;
+ std::vector<int64_t> bad_rows_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_;
+ std::atomic<int64_t> rows_generated_;
+ float scale_factor_;
+ int64_t batch_size_;
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class PartGenerator : public TpchTableGenerator {
+ public:
+ explicit PartGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetPartOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch, gen_->NextPartBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->part_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class PartSuppGenerator : public TpchTableGenerator {
+ public:
+ explicit PartSuppGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetPartSuppOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextPartSuppBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->partsupp_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class CustomerGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, gen_list_));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) tld.rng.seed(seq);
+
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(C_CUSTKEY) \
+ F(C_NAME) \
+ F(C_ADDRESS) \
+ F(C_NATIONKEY) \
+ F(C_PHONE) \
+ F(C_ACCTBAL) \
+ F(C_MKTSEGMENT) \
+ F(C_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct CUSTOMER {
+ enum {
+ FOR_EACH_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) {#col, CUSTOMER::col},
+ const std::unordered_map<std::string, int> kNameMap = {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kGenerators = {FOR_EACH_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(),
+ utf8(),
+ utf8(),
+ int32(),
+ fixed_size_binary(15),
+ decimal(12, 2),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index) {
+ if (done_.load()) return Status::OK();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.custkey_start = rows_generated_.fetch_add(batch_size_);
+ if (tld.custkey_start >= rows_to_generate_) return Status::OK();
+
+ tld.to_generate = std::min(batch_size_, rows_to_generate_ - tld.custkey_start);
+
+ tld.batch.resize(CUSTOMER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for (int col : gen_list_) RETURN_NOT_OK(kGenerators[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for (size_t i = 0; i < gen_list_.size(); i++) {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_generated_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_generated_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if (is_last_batch) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ return Status::OK();
+ }
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(kTypes[column], tld.to_generate, {nullptr, std::move(buff)});
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status C_CUSTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_CUSTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_CUSTKEY));
+ int32_t* c_custkey = reinterpret_cast<int32_t*>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ c_custkey[irow] = static_cast<int32_t>(tld.custkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_NAME].kind() == Datum::NONE) {
+ RETURN_NOT_OK(C_CUSTKEY(thread_index));
+ const int32_t* c_custkey = reinterpret_cast<const int32_t*>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->data());
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ const char* customer = "Customer#";
+ const size_t customer_length = std::strlen(customer);
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ int num_digits = GetNumDigits(c_custkey[irow]);
+ int num_chars = std::max(num_digits, 9);
+ offsets[irow + 1] =
+ static_cast<int32_t>(offsets[irow] + num_chars + customer_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[tld.to_generate]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = str + offsets[irow];
+ std::memcpy(out, customer, customer_length);
+ AppendNumberPaddedToNineDigits(out + customer_length, c_custkey[irow]);
+ }
+ ArrayData ad(utf8(), tld.to_generate,
+ {nullptr, std::move(offset_buff), std::move(str_buff)});
+ tld.batch[CUSTOMER::C_NAME] = std::move(ad);
+ }
+ return Status::OK();
+ }
+
+ Status C_ADDRESS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_ADDRESS].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[CUSTOMER::C_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status C_NATIONKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_NATIONKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t* c_nationkey = reinterpret_cast<int32_t*>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status C_PHONE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_PHONE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(C_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[CUSTOMER::C_PHONE]);
+ const int32_t* c_nationkey = reinterpret_cast<const int32_t*>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->data());
+ char* c_phone = reinterpret_cast<char*>(
+ tld.batch[CUSTOMER::C_PHONE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ GeneratePhoneNumber(c_phone + irow * byte_width, tld.rng, c_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_ACCTBAL(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_ACCTBAL].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_ACCTBAL));
+ Decimal128* c_acctbal = reinterpret_cast<Decimal128*>(
+ tld.batch[CUSTOMER::C_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_acctbal[irow] = {dist(tld.rng)};
+ }
+ return Status::OK();
+ }
+
+ Status C_MKTSEGMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_MKTSEGMENT].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_MKTSEGMENT));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[CUSTOMER::C_MKTSEGMENT]);
+ char* c_mktsegment = reinterpret_cast<char*>(
+ tld.batch[CUSTOMER::C_MKTSEGMENT].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int32_t> dist(0, kNumSegments - 1);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = c_mktsegment + irow * byte_width;
+ int str_idx = dist(tld.rng);
+ std::strncpy(out, Segments[str_idx], byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[CUSTOMER::C_COMMENT],
+ g_text.GenerateComments(tld.to_generate, 29, 116, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ random::pcg32_fast rng;
+ int64_t custkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_{0};
+ std::atomic<int64_t> rows_generated_ = {0};
+ float scale_factor_{0};
+ int64_t batch_size_{0};
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class OrdersGenerator : public TpchTableGenerator {
+ public:
+ explicit OrdersGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetOrdersOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch, gen_->NextOrdersBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->orders_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class LineitemGenerator : public TpchTableGenerator {
+ public:
+ explicit LineitemGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetLineItemOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextLineItemBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->lineitem_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class NationGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, float /*scale_factor*/,
+ int64_t /*batch_size*/) override {
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, column_indices_));
+ rng_.seed(arrow_vendored::pcg_extras::seed_seq_from<std::random_device>{});
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t /*num_threads*/, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback /*schedule_task_callback*/) override {
+ std::shared_ptr<Buffer> N_NATIONKEY_buffer =
+ Buffer::Wrap(kNationKey, sizeof(kNationKey));
+ ArrayData N_NATIONKEY_arraydata(int32(), kRowCount,
+ {nullptr, std::move(N_NATIONKEY_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> N_NAME_buffer,
+ AllocateBuffer(kRowCount * kNameByteWidth));
+ char* N_NAME = reinterpret_cast<char*>(N_NAME_buffer->mutable_data());
+ for (size_t i = 0; i < kRowCount; i++)
+ std::strncpy(N_NAME + kNameByteWidth * i, kCountryNames[i], kNameByteWidth);
+ ArrayData N_NAME_arraydata(fixed_size_binary(kNameByteWidth), kRowCount,
+ {nullptr, std::move(N_NAME_buffer)});
+
+ std::shared_ptr<Buffer> N_REGIONKEY_buffer =
+ Buffer::Wrap(kRegionKey, sizeof(kRegionKey));
+ ArrayData N_REGIONKEY_arraydata(int32(), kRowCount,
+ {nullptr, std::move(N_REGIONKEY_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(Datum N_COMMENT_datum,
+ g_text.GenerateComments(kRowCount, 31, 114, rng_));
+
+ std::vector<Datum> fields = {
+ std::move(N_NATIONKEY_arraydata), std::move(N_NAME_arraydata),
+ std::move(N_REGIONKEY_arraydata), std::move(N_COMMENT_datum)};
+
+ std::vector<Datum> result;
+ for (const int& col : column_indices_) result.push_back(fields[col]);
+ ARROW_ASSIGN_OR_RAISE(ExecBatch batch, ExecBatch::Make(std::move(result)));
+ output_callback(std::move(batch));
+ finished_callback(static_cast<int64_t>(1));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ random::pcg32_fast rng_;
+
+ static constexpr size_t kRowCount = 25;
+ static constexpr int32_t kNameByteWidth = 25;
+ const int32_t kNationKey[kRowCount] = {0, 1, 2, 3, 4, 5, 6, 7, 8,
+ 9, 10, 11, 12, 13, 14, 15, 16, 17,
+ 18, 19, 20, 21, 22, 23, 24};
+ const char* kCountryNames[kRowCount] = {
+ "ALGERIA", "ARGENTINA", "BRAZIL", "CANADA", "EGYPT",
+ "ETHIOPIA", "FRANCE", "GERMANY", "INDIA", "INDONESIA",
+ "IRAN", "IRAQ", "JAPAN", "JORDAN", "KENYA",
+ "MOROCCO", "MOZAMBIQUE", "PERU", "CHINA", "ROMANIA",
+ "SAUDI ARABIA", "VIETNAM", "RUSSIA", "UNITED KINGDOM", "UNITED STATES"};
+ const int32_t kRegionKey[kRowCount] = {0, 1, 1, 1, 4, 0, 3, 3, 2, 2, 4, 4, 2,
+ 4, 0, 0, 0, 1, 2, 3, 4, 2, 3, 3, 1};
+
+ struct NATION {
+ enum {
Review comment:
You can instead write:
```suggestion
enum class NATION {
```
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
Review comment:
Can you instead use e.g. a `std::vector<std::string>`? We're trying to write idiomatic C++ code as much as possible.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
Review comment:
Just for the record, TPC-H specifies that phone numbers never have a leading zero?
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
Review comment:
I may be mistaken (this code is difficult to follow and you didn't provide any comments), but it looks like this is doing home-grown thread synchronization. Is there a strong reason not to use the existing `ThreadPool` facility?
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
Review comment:
Again, this could probably use one of the formatting facilities.
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
Review comment:
Also, note you can use a verification class/struct that would expose the VerifyXXX methods and accumulate statistics about the data in member attributes. This would probably be more idiomatic.
##########
File path: cpp/src/arrow/compute/kernels/vector_sort.cc
##########
@@ -893,7 +893,8 @@ class TableSorter {
TableSorter(ExecContext* ctx, uint64_t* indices_begin, uint64_t* indices_end,
const Table& table, const SortOptions& options)
- : ctx_(ctx),
Review comment:
@save-buffer Do you remember where the "resetting the error" part happened?
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
Review comment:
It would be better to call `GetRandomSeed` (see `arrow/util/io_util.h`).
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ if (!partsupp_output_queue_.empty()) {
+ ExecBatch result = std::move(partsupp_output_queue_.front());
+ partsupp_output_queue_.pop();
+ return std::move(result);
+ }
+ }
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+ if (!part_cols_.empty()) {
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch part_batch,
+ ExecBatch::Make(std::move(part_result)));
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ part_output_queue_.emplace(std::move(part_batch));
+ }
+ }
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ for (size_t i = 1; i < partsupp_results.size(); i++)
+ partsupp_output_queue_.emplace(std::move(partsupp_results[i]));
+ }
+ return std::move(partsupp_results[0]);
+ }
+
+ private:
+#define FOR_EACH_PART_COLUMN(F) \
+ F(P_PARTKEY) \
+ F(P_NAME) \
+ F(P_MFGR) \
+ F(P_BRAND) \
+ F(P_TYPE) \
+ F(P_SIZE) \
+ F(P_CONTAINER) \
+ F(P_RETAILPRICE) \
+ F(P_COMMENT)
+
+#define FOR_EACH_PARTSUPP_COLUMN(F) \
+ F(PS_PARTKEY) \
+ F(PS_SUPPKEY) \
+ F(PS_AVAILQTY) \
+ F(PS_SUPPLYCOST) \
+ F(PS_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct PART {
+ enum {
+ FOR_EACH_PART_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct PARTSUPP {
+ enum {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, PART::col},
+ const std::unordered_map<std::string, int> kPartNameMap = {
+ FOR_EACH_PART_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, PARTSUPP::col},
+ const std::unordered_map<std::string, int> kPartsuppNameMap = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kPartGenerators = {FOR_EACH_PART_COLUMN(MAKE_FN_ARRAY)};
+ std::vector<GenerateColumnFn> kPartsuppGenerators = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kPartTypes = {
+ int32(), utf8(), fixed_size_binary(25), fixed_size_binary(10),
+ utf8(), int32(), fixed_size_binary(10), decimal(12, 2),
+ utf8(),
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kPartsuppTypes = {
+ int32(), int32(), int32(), decimal(12, 2), utf8(),
+ };
+
+ Status AllocatePartBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.part[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.part_to_generate * byte_width));
+ ArrayData ad(kPartTypes[column], tld.part_to_generate, {nullptr, std::move(buff)});
+ tld.part[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status P_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_PARTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_PARTKEY));
+ int32_t* p_partkey = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) {
+ p_partkey[i] = static_cast<int32_t>(tld.partkey_start + i + 1);
+ ARROW_DCHECK(1 <= p_partkey[i] && p_partkey[i] <= part_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_NAME].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(
+ 0, static_cast<uint8_t>(kNumNameParts - 1));
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dist(tld.rng));
+ tld.string_indices[irow * 5 + ipart] = name_part_index;
+ string_length += std::strlen(NameParts[name_part_index]);
+ }
+ // Add 4 because there is a space between each word (i.e. four spaces)
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length + 4);
+ }
+ // Add an extra byte for the space after in the very last string.
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate] + 1));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 5 + ipart];
+ const char* part = NameParts[name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ *row++ = ' ';
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_NAME], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_NAME] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_MFGR(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_MFGR].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* manufacturer = "Manufacturer#";
+ const size_t manufacturer_length = std::strlen(manufacturer);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_MFGR));
+ char* p_mfgr = reinterpret_cast<char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ std::strncpy(p_mfgr + byte_width * irow, manufacturer, byte_width);
+ char mfgr_id = '0' + dist(tld.rng);
+ *(p_mfgr + byte_width * irow + manufacturer_length) = mfgr_id;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_BRAND(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_BRAND].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_MFGR(thread_index));
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* brand = "Brand#";
+ const size_t brand_length = std::strlen(brand);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_BRAND));
+ const char* p_mfgr = reinterpret_cast<const char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->data());
+ char* p_brand = reinterpret_cast<char*>(
+ tld.part[PART::P_BRAND].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_BRAND]);
+ int32_t mfgr_byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ const size_t mfgr_id_offset = std::strlen("Manufacturer#");
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = p_brand + byte_width * irow;
+ char mfgr_id = *(p_mfgr + irow * mfgr_byte_width + mfgr_id_offset);
+ char brand_id = '0' + dist(tld.rng);
+ std::strncpy(row, brand, byte_width);
+ *(row + brand_length) = mfgr_id;
+ *(row + brand_length + 1) = brand_id;
+ irow += 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_TYPE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_TYPE].kind() == Datum::NONE) {
+ using D = std::uniform_int_distribution<uint32_t>;
+ D dists[] = {
+ D{0, static_cast<uint8_t>(kNumTypes_1 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_2 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_3 - 1)},
+ };
+
+ const char** types[] = {Types_1, Types_2, Types_3};
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dists[ipart](tld.rng));
+ tld.string_indices[irow * 3 + ipart] = name_part_index;
+ string_length += std::strlen(types[ipart][name_part_index]);
+ }
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate]));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 3 + ipart];
+ const char* part = types[ipart][name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_TYPE], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_TYPE] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_SIZE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_SIZE].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(1, 50);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_SIZE));
+ int32_t* p_size = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_SIZE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) p_size[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status P_CONTAINER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_CONTAINER].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist1(
+ 0, static_cast<uint8_t>(kNumContainers_1 - 1));
+ std::uniform_int_distribution<int> dist2(
+ 0, static_cast<uint8_t>(kNumContainers_2 - 1));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_CONTAINER));
+ char* p_container = reinterpret_cast<char*>(
+ tld.part[PART::P_CONTAINER].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_CONTAINER]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int container1_idx = dist1(tld.rng);
+ int container2_idx = dist2(tld.rng);
+ const char* container1 = Containers_1[container1_idx];
+ const char* container2 = Containers_2[container2_idx];
+ size_t container1_length = std::strlen(container1);
+ size_t container2_length = std::strlen(container2);
+
+ char* row = p_container + byte_width * irow;
+ std::strncpy(row, container1, byte_width);
+ std::memcpy(row + container1_length, container2, container2_length);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_RETAILPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_RETAILPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_RETAILPRICE));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+ Decimal128* p_retailprice = reinterpret_cast<Decimal128*>(
+ tld.part[PART::P_RETAILPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int32_t partkey = p_partkey[irow];
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ p_retailprice[irow] = {retail_price};
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.part[PART::P_COMMENT],
+ g_text.GenerateComments(tld.part_to_generate, 5, 22, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ int64_t PartsuppBatchesToGenerate(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t num_batches = (ps_to_generate + batch_size_ - 1) / batch_size_;
+ return num_batches;
+ }
+
+ Status InitPartsupp(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.generated_partsupp.reset();
+ int64_t num_batches = PartsuppBatchesToGenerate(thread_index);
+ tld.partsupp.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.partsupp) {
+ batch.resize(PARTSUPP::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocatePartSuppBatch(size_t thread_index, size_t ibatch, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartsuppTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kPartsuppTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.partsupp[ibatch][column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status PS_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_PARTKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_PARTKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_PARTKEY));
+ int32_t* ps_partkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run; ipartsupp++, irun++)
+ ps_partkey[batch_offset++] = p_partkey[ipart];
+
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPKEY));
+ int32_t* ps_suppkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run;
+ ipartsupp++, irun++) {
+ int32_t supplier = static_cast<int32_t>(ipartsupp);
+ int32_t partkey = p_partkey[ipart];
+ ps_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
Review comment:
What is this? Some magic formula from TPC-H? Can you please add a comment?
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ if (!partsupp_output_queue_.empty()) {
+ ExecBatch result = std::move(partsupp_output_queue_.front());
+ partsupp_output_queue_.pop();
+ return std::move(result);
+ }
+ }
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+ if (!part_cols_.empty()) {
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch part_batch,
+ ExecBatch::Make(std::move(part_result)));
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ part_output_queue_.emplace(std::move(part_batch));
+ }
+ }
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ for (size_t i = 1; i < partsupp_results.size(); i++)
+ partsupp_output_queue_.emplace(std::move(partsupp_results[i]));
+ }
+ return std::move(partsupp_results[0]);
+ }
+
+ private:
+#define FOR_EACH_PART_COLUMN(F) \
+ F(P_PARTKEY) \
+ F(P_NAME) \
+ F(P_MFGR) \
+ F(P_BRAND) \
+ F(P_TYPE) \
+ F(P_SIZE) \
+ F(P_CONTAINER) \
+ F(P_RETAILPRICE) \
+ F(P_COMMENT)
+
+#define FOR_EACH_PARTSUPP_COLUMN(F) \
+ F(PS_PARTKEY) \
+ F(PS_SUPPKEY) \
+ F(PS_AVAILQTY) \
+ F(PS_SUPPLYCOST) \
+ F(PS_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct PART {
+ enum {
+ FOR_EACH_PART_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct PARTSUPP {
+ enum {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, PART::col},
+ const std::unordered_map<std::string, int> kPartNameMap = {
+ FOR_EACH_PART_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, PARTSUPP::col},
+ const std::unordered_map<std::string, int> kPartsuppNameMap = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kPartGenerators = {FOR_EACH_PART_COLUMN(MAKE_FN_ARRAY)};
+ std::vector<GenerateColumnFn> kPartsuppGenerators = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kPartTypes = {
+ int32(), utf8(), fixed_size_binary(25), fixed_size_binary(10),
+ utf8(), int32(), fixed_size_binary(10), decimal(12, 2),
+ utf8(),
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kPartsuppTypes = {
+ int32(), int32(), int32(), decimal(12, 2), utf8(),
+ };
+
+ Status AllocatePartBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.part[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.part_to_generate * byte_width));
+ ArrayData ad(kPartTypes[column], tld.part_to_generate, {nullptr, std::move(buff)});
+ tld.part[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status P_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_PARTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_PARTKEY));
+ int32_t* p_partkey = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) {
+ p_partkey[i] = static_cast<int32_t>(tld.partkey_start + i + 1);
+ ARROW_DCHECK(1 <= p_partkey[i] && p_partkey[i] <= part_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_NAME].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(
+ 0, static_cast<uint8_t>(kNumNameParts - 1));
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dist(tld.rng));
+ tld.string_indices[irow * 5 + ipart] = name_part_index;
+ string_length += std::strlen(NameParts[name_part_index]);
+ }
+ // Add 4 because there is a space between each word (i.e. four spaces)
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length + 4);
+ }
+ // Add an extra byte for the space after in the very last string.
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate] + 1));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 5 + ipart];
+ const char* part = NameParts[name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ *row++ = ' ';
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_NAME], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_NAME] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_MFGR(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_MFGR].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* manufacturer = "Manufacturer#";
+ const size_t manufacturer_length = std::strlen(manufacturer);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_MFGR));
+ char* p_mfgr = reinterpret_cast<char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ std::strncpy(p_mfgr + byte_width * irow, manufacturer, byte_width);
+ char mfgr_id = '0' + dist(tld.rng);
+ *(p_mfgr + byte_width * irow + manufacturer_length) = mfgr_id;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_BRAND(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_BRAND].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_MFGR(thread_index));
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* brand = "Brand#";
+ const size_t brand_length = std::strlen(brand);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_BRAND));
+ const char* p_mfgr = reinterpret_cast<const char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->data());
+ char* p_brand = reinterpret_cast<char*>(
+ tld.part[PART::P_BRAND].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_BRAND]);
+ int32_t mfgr_byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ const size_t mfgr_id_offset = std::strlen("Manufacturer#");
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = p_brand + byte_width * irow;
+ char mfgr_id = *(p_mfgr + irow * mfgr_byte_width + mfgr_id_offset);
+ char brand_id = '0' + dist(tld.rng);
+ std::strncpy(row, brand, byte_width);
+ *(row + brand_length) = mfgr_id;
+ *(row + brand_length + 1) = brand_id;
+ irow += 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_TYPE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_TYPE].kind() == Datum::NONE) {
+ using D = std::uniform_int_distribution<uint32_t>;
+ D dists[] = {
+ D{0, static_cast<uint8_t>(kNumTypes_1 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_2 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_3 - 1)},
+ };
+
+ const char** types[] = {Types_1, Types_2, Types_3};
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dists[ipart](tld.rng));
+ tld.string_indices[irow * 3 + ipart] = name_part_index;
+ string_length += std::strlen(types[ipart][name_part_index]);
+ }
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate]));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 3 + ipart];
+ const char* part = types[ipart][name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_TYPE], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_TYPE] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_SIZE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_SIZE].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(1, 50);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_SIZE));
+ int32_t* p_size = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_SIZE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) p_size[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status P_CONTAINER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_CONTAINER].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist1(
+ 0, static_cast<uint8_t>(kNumContainers_1 - 1));
+ std::uniform_int_distribution<int> dist2(
+ 0, static_cast<uint8_t>(kNumContainers_2 - 1));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_CONTAINER));
+ char* p_container = reinterpret_cast<char*>(
+ tld.part[PART::P_CONTAINER].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_CONTAINER]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int container1_idx = dist1(tld.rng);
+ int container2_idx = dist2(tld.rng);
+ const char* container1 = Containers_1[container1_idx];
+ const char* container2 = Containers_2[container2_idx];
+ size_t container1_length = std::strlen(container1);
+ size_t container2_length = std::strlen(container2);
+
+ char* row = p_container + byte_width * irow;
+ std::strncpy(row, container1, byte_width);
+ std::memcpy(row + container1_length, container2, container2_length);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_RETAILPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_RETAILPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_RETAILPRICE));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+ Decimal128* p_retailprice = reinterpret_cast<Decimal128*>(
+ tld.part[PART::P_RETAILPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int32_t partkey = p_partkey[irow];
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ p_retailprice[irow] = {retail_price};
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.part[PART::P_COMMENT],
+ g_text.GenerateComments(tld.part_to_generate, 5, 22, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ int64_t PartsuppBatchesToGenerate(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t num_batches = (ps_to_generate + batch_size_ - 1) / batch_size_;
+ return num_batches;
+ }
+
+ Status InitPartsupp(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.generated_partsupp.reset();
+ int64_t num_batches = PartsuppBatchesToGenerate(thread_index);
+ tld.partsupp.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.partsupp) {
+ batch.resize(PARTSUPP::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocatePartSuppBatch(size_t thread_index, size_t ibatch, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartsuppTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kPartsuppTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.partsupp[ibatch][column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status PS_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_PARTKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_PARTKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_PARTKEY));
+ int32_t* ps_partkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run; ipartsupp++, irun++)
+ ps_partkey[batch_offset++] = p_partkey[ipart];
+
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPKEY));
+ int32_t* ps_suppkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run;
+ ipartsupp++, irun++) {
+ int32_t supplier = static_cast<int32_t>(ipartsupp);
+ int32_t partkey = p_partkey[ipart];
+ ps_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_AVAILQTY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_AVAILQTY]) {
+ tld.generated_partsupp[PARTSUPP::PS_AVAILQTY] = true;
+ std::uniform_int_distribution<int32_t> dist(1, 9999);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_AVAILQTY));
+ int32_t* ps_availqty =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for (int64_t irun = 0; irun < next_run; irun++) ps_availqty[irun] = dist(tld.rng);
+
+ tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPLYCOST(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST] = true;
+ std::uniform_int_distribution<int64_t> dist(100, 100000);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(
+ AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPLYCOST));
+ Decimal128* ps_supplycost =
+ reinterpret_cast<Decimal128*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for (int64_t irun = 0; irun < next_run; irun++)
+ ps_supplycost[irun] = {dist(tld.rng)};
+
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PARTSUPP::PS_COMMENT].kind() == Datum::NONE) {
+ int64_t irow = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ int64_t num_rows = std::min(batch_size_, ps_to_generate - irow);
+ ARROW_ASSIGN_OR_RAISE(tld.partsupp[ibatch][PARTSUPP::PS_COMMENT],
+ g_text.GenerateComments(num_rows, 49, 198, tld.rng));
+ irow += num_rows;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ std::vector<Datum> part;
+ std::vector<int8_t> string_indices;
+ int64_t part_to_generate{0};
+ int64_t partkey_start{0};
+
+ std::vector<std::vector<Datum>> partsupp;
+ std::bitset<PARTSUPP::kNumCols> generated_partsupp;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex part_output_queue_mutex_;
+ std::mutex partsupp_output_queue_mutex_;
+ std::queue<ExecBatch> part_output_queue_;
+ std::queue<ExecBatch> partsupp_output_queue_;
+ int64_t batch_size_{0};
+ float scale_factor_{0};
+ int64_t part_rows_to_generate_{0};
+ int64_t part_rows_generated_{0};
+ std::vector<int> part_cols_;
+ std::vector<int> partsupp_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<int64_t> part_batches_generated_ = {0};
+ std::atomic<int64_t> partsupp_batches_generated_ = {0};
+ static constexpr int64_t kPartSuppRowsPerPart = 4;
+};
+
+class OrdersAndLineItemGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ tld.items_per_order.resize(batch_size_);
+ tld.rng.seed(seq);
+ }
+ orders_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000 * 10);
+ }
+ return Status::OK();
+ }
+
+ int64_t orders_batches_generated() const { return orders_batches_generated_.load(); }
+
+ int64_t lineitem_batches_generated() const {
+ return lineitem_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetOrdersOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kOrdersTypes, kOrdersNameMap, orders_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetLineItemOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kLineitemTypes, kLineitemNameMap, lineitem_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextOrdersBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if (!orders_output_queue_.empty()) {
+ ExecBatch batch = std::move(orders_output_queue_.front());
+ orders_output_queue_.pop();
+ return std::move(batch);
+ } else if (orders_rows_generated_ == orders_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate =
+ std::min(batch_size_, orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.first_batch_offset = 0;
+ tld.generated_lineitem.reset();
+
+ for (int col : orders_cols_) RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for (int col : lineitem_cols_) RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for (size_t i = 0; i < orders_cols_.size(); i++) {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ if (!lineitem_cols_.empty()) {
+ std::vector<ExecBatch> lineitem_results;
+ for (size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++) {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for (size_t icol = 0; icol < lineitem_cols_.size(); icol++) {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(lineitem_output_queue_mutex_);
+ for (ExecBatch& eb : lineitem_results) {
+ lineitem_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(orders_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextLineItemBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ExecBatch queued;
+ bool from_queue = false;
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ if (!lineitem_output_queue_.empty()) {
+ queued = std::move(lineitem_output_queue_.front());
+ lineitem_output_queue_.pop();
+ from_queue = true;
+ }
+ }
+ tld.first_batch_offset = 0;
+ if (from_queue) {
+ ARROW_DCHECK(queued.length <= batch_size_);
+ tld.first_batch_offset = queued.length;
+ if (queued.length == batch_size_) return std::move(queued);
+ }
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if (orders_rows_generated_ == orders_rows_to_generate_) {
+ if (from_queue) return std::move(queued);
+ return util::nullopt;
+ }
+
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate =
+ std::min(batch_size_, orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1ll);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.generated_lineitem.reset();
+ if (from_queue) {
+ lineitem_batches_generated_.fetch_sub(1);
+ for (size_t i = 0; i < lineitem_cols_.size(); i++)
+ if (tld.lineitem[0][lineitem_cols_[i]].kind() == Datum::NONE)
+ tld.lineitem[0][lineitem_cols_[i]] = std::move(queued[i]);
+ }
+
+ for (int col : orders_cols_) RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for (int col : lineitem_cols_) RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ if (!orders_cols_.empty()) {
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for (size_t i = 0; i < orders_cols_.size(); i++) {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch orders_batch,
+ ExecBatch::Make(std::move(orders_result)));
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ orders_output_queue_.emplace(std::move(orders_batch));
+ }
+ }
+ std::vector<ExecBatch> lineitem_results;
+ for (size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++) {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for (size_t icol = 0; icol < lineitem_cols_.size(); icol++) {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ lineitem_batches_generated_.fetch_add(static_cast<int64_t>(lineitem_results.size()));
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ for (size_t i = 1; i < lineitem_results.size(); i++)
+ lineitem_output_queue_.emplace(std::move(lineitem_results[i]));
+ }
+ return std::move(lineitem_results[0]);
+ }
+
+ private:
+#define FOR_EACH_ORDERS_COLUMN(F) \
+ F(O_ORDERKEY) \
+ F(O_CUSTKEY) \
+ F(O_ORDERSTATUS) \
+ F(O_TOTALPRICE) \
+ F(O_ORDERDATE) \
+ F(O_ORDERPRIORITY) \
+ F(O_CLERK) \
+ F(O_SHIPPRIORITY) \
+ F(O_COMMENT)
+
+#define FOR_EACH_LINEITEM_COLUMN(F) \
+ F(L_ORDERKEY) \
+ F(L_PARTKEY) \
+ F(L_SUPPKEY) \
+ F(L_LINENUMBER) \
+ F(L_QUANTITY) \
+ F(L_EXTENDEDPRICE) \
+ F(L_DISCOUNT) \
+ F(L_TAX) \
+ F(L_RETURNFLAG) \
+ F(L_LINESTATUS) \
+ F(L_SHIPDATE) \
+ F(L_COMMITDATE) \
+ F(L_RECEIPTDATE) \
+ F(L_SHIPINSTRUCT) \
+ F(L_SHIPMODE) \
+ F(L_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct ORDERS {
+ enum {
+ FOR_EACH_ORDERS_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct LINEITEM {
+ enum {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, ORDERS::col},
+ const std::unordered_map<std::string, int> kOrdersNameMap = {
+ FOR_EACH_ORDERS_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, LINEITEM::col},
+ const std::unordered_map<std::string, int> kLineitemNameMap = {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ const std::vector<GenerateColumnFn> kOrdersGenerators = {
+ FOR_EACH_ORDERS_COLUMN(MAKE_FN_ARRAY)};
+ const std::vector<GenerateColumnFn> kLineitemGenerators = {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kOrdersTypes = {int32(),
+ int32(),
+ fixed_size_binary(1),
+ decimal(12, 2),
+ date32(),
+ fixed_size_binary(15),
+ fixed_size_binary(15),
+ int32(),
+ utf8()};
+
+ const std::vector<std::shared_ptr<DataType>> kLineitemTypes = {
+ int32(),
+ int32(),
+ int32(),
+ int32(),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ fixed_size_binary(1),
+ fixed_size_binary(1),
+ date32(),
+ date32(),
+ date32(),
+ fixed_size_binary(25),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status AllocateOrdersBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.orders[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.orders_to_generate * byte_width));
+ ArrayData ad(kOrdersTypes[column], tld.orders_to_generate,
+ {nullptr, std::move(buff)});
+ tld.orders[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status O_ORDERKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERKEY));
+ int32_t* o_orderkey = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ int32_t orderkey_index = static_cast<int32_t>(tld.orderkey_start + i);
+ int32_t index_of_run = orderkey_index / 8;
+ int32_t index_in_run = orderkey_index % 8;
+ o_orderkey[i] = (index_of_run * 32 + index_in_run + 1);
+ ARROW_DCHECK(1 <= o_orderkey[i] && o_orderkey[i] <= 4 * orders_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CUSTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_CUSTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CUSTKEY));
+
+ // Spec says it must be a random number between 1 and SF*150000 that is not
+ // divisible by 3. Rather than repeatedly generating numbers until we get to
+ // a non-divisible-by-3 number, we just generate a number between
+ // 0 and SF * 50000 - 1, multiply by 3, and then add either 1 or 2.
+ int32_t sf_50k = static_cast<int32_t>(scale_factor_ * 50000);
+ std::uniform_int_distribution<int32_t> base_dist(0, sf_50k - 1);
+ std::uniform_int_distribution<int32_t> offset_dist(1, 2);
+ int32_t* o_custkey = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_CUSTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++)
+ o_custkey[i] = 3 * base_dist(tld.rng) + offset_dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERSTATUS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERSTATUS].kind() == Datum::NONE) {
+ RETURN_NOT_OK(L_LINESTATUS(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERSTATUS));
+
+ char* o_orderstatus = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_ORDERSTATUS].array()->buffers[1]->mutable_data());
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ bool all_f = true;
+ bool all_o = true;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ const char* l_linestatus = reinterpret_cast<const char*>(
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->buffers[1]->data());
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++, batch_offset++) {
+ all_f &= l_linestatus[batch_offset] == 'F';
+ all_o &= l_linestatus[batch_offset] == 'O';
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ ARROW_DCHECK(!(all_f && all_o));
+ if (all_f)
+ o_orderstatus[iorder] = 'F';
+ else if (all_o)
+ o_orderstatus[iorder] = 'O';
+ else
+ o_orderstatus[iorder] = 'P';
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_TOTALPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_TOTALPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(L_EXTENDEDPRICE(thread_index));
+ RETURN_NOT_OK(L_TAX(thread_index));
+ RETURN_NOT_OK(L_DISCOUNT(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_TOTALPRICE));
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ int64_t sum = 0;
+ Decimal128* o_totalprice = reinterpret_cast<Decimal128*>(
+ tld.orders[ORDERS::O_TOTALPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ const Decimal128* l_extendedprice = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->buffers[1]->data());
+ const Decimal128* l_tax = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->data());
+ const Decimal128* l_discount = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->buffers[1]->data());
+
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++, batch_offset++) {
+ int64_t eprice = static_cast<int64_t>(l_extendedprice[batch_offset]);
+ int64_t tax = static_cast<int64_t>(l_tax[batch_offset]);
+ int64_t discount = static_cast<int64_t>(l_discount[batch_offset]);
+ sum += (eprice * (100 + tax) * (100 - discount));
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ sum /= 100 * 100;
+ o_totalprice[iorder] = {sum};
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERDATE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERDATE));
+
+ std::uniform_int_distribution<uint32_t> dist(kStartDate, kEndDate - 151);
+ uint32_t* o_orderdate = reinterpret_cast<uint32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) o_orderdate[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERPRIORITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERPRIORITY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERPRIORITY));
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_ORDERPRIORITY]);
+ std::uniform_int_distribution<int32_t> dist(0, kNumPriorities - 1);
+ char* o_orderpriority = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_ORDERPRIORITY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ const char* str = Priorities[dist(tld.rng)];
+ std::strncpy(o_orderpriority + i * byte_width, str, byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CLERK(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_CLERK].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CLERK));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_CLERK]);
+ int64_t max_clerk_id = static_cast<int64_t>(scale_factor_ * 1000);
+ std::uniform_int_distribution<int64_t> dist(1, max_clerk_id);
+ char* o_clerk = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_CLERK].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ const char* clerk = "Clerk#";
+ const size_t clerk_length = std::strlen(clerk);
+ int64_t clerk_number = dist(tld.rng);
+ char* output = o_clerk + i * byte_width;
+ std::strncpy(output, clerk, byte_width);
+ AppendNumberPaddedToNineDigits(output + clerk_length, clerk_number);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_SHIPPRIORITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_SHIPPRIORITY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_SHIPPRIORITY));
+ int32_t* o_shippriority = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_SHIPPRIORITY].array()->buffers[1]->mutable_data());
+ std::memset(o_shippriority, 0, tld.orders_to_generate * sizeof(int32_t));
+ }
+ return Status::OK();
+ }
+
+ Status O_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.orders[ORDERS::O_COMMENT],
+ g_text.GenerateComments(tld.orders_to_generate, 19, 78, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ Status GenerateRowCounts(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ std::uniform_int_distribution<int> length_dist(1, 7);
+ tld.lineitem_to_generate = 0;
+ tld.items_per_order.clear();
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ int length = length_dist(tld.rng);
+ tld.items_per_order.push_back(length);
+ tld.lineitem_to_generate += length;
+ }
+ int64_t num_batches =
+ (tld.first_batch_offset + tld.lineitem_to_generate + batch_size_ - 1) /
+ batch_size_;
+ tld.lineitem.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.lineitem) {
+ batch.resize(LINEITEM::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocateLineItemBufferIfNeeded(size_t thread_index, size_t ibatch, int column,
+ size_t& out_batch_offset) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.lineitem[ibatch][column].kind() == Datum::NONE) {
+ int32_t byte_width = arrow::internal::GetByteWidth(*kLineitemTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kLineitemTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.lineitem[ibatch][column] = std::move(ad);
+ out_batch_offset = 0;
+ }
+ if (ibatch == 0) out_batch_offset = tld.first_batch_offset;
+
+ return Status::OK();
+ }
+
+ Status L_ORDERKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_ORDERKEY]) {
+ tld.generated_lineitem[LINEITEM::L_ORDERKEY] = true;
+ RETURN_NOT_OK(O_ORDERKEY(thread_index));
+ const int32_t* o_orderkey = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_ORDERKEY, batch_offset));
+ int32_t* l_linenumber =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_ORDERKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_linenumber[batch_offset++] = o_orderkey[iorder];
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_ORDERKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_PARTKEY]) {
+ tld.generated_lineitem[LINEITEM::L_PARTKEY] = true;
+
+ size_t ibatch = 0;
+ int32_t max_partkey = static_cast<int32_t>(scale_factor_ * 200000);
+ std::uniform_int_distribution<int32_t> dist(1, max_partkey);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_PARTKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_partkey =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_partkey[batch_offset] = dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SUPPKEY]) {
+ tld.generated_lineitem[LINEITEM::L_SUPPKEY] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 3);
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset = 0;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SUPPKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_suppkey =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const int32_t* l_partkey = reinterpret_cast<const int32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ for (int64_t i = 0; i < next_run; i++) {
+ int32_t supplier = dist(tld.rng);
+ int32_t partkey = l_partkey[batch_offset];
+ // Fun fact: the parentheses for this expression are unbalanced in the TPC-H
+ // spec.
+ l_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SUPPKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINENUMBER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_LINENUMBER]) {
+ tld.generated_lineitem[LINEITEM::L_LINENUMBER] = true;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_LINENUMBER, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_linenumber =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_LINENUMBER]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++) {
+ l_linenumber[batch_offset++] = (iline + 1);
+ ARROW_DCHECK(1 <= (iline + 1) && (iline + 1) <= 7);
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINENUMBER].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_QUANTITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_QUANTITY]) {
+ tld.generated_lineitem[LINEITEM::L_QUANTITY] = true;
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int64_t> dist(1, 50);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_QUANTITY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_quantity =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_QUANTITY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t i = 0; i < next_run; i++) {
+ // Multiply by 100 because the type is decimal(12, 2), so the decimal goes after
+ // two digits
+ int64_t quantity = dist(tld.rng) * 100;
+ l_quantity[batch_offset++] = {quantity};
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_EXTENDEDPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE]) {
+ tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+ RETURN_NOT_OK(L_QUANTITY(thread_index));
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_EXTENDEDPRICE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ const int32_t* l_partkey = reinterpret_cast<const int32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ const Decimal128* l_quantity = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->buffers[1]->data());
+ Decimal128* l_extendedprice =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ int64_t partkey = static_cast<int64_t>(l_partkey[batch_offset]);
+ // Divide by 100 to recover the integer representation (not Decimal).
+ int64_t quantity = static_cast<int64_t>(l_quantity[batch_offset]) / 100;
+
+ // Spec says to divide by 100, but that happens automatically due to this being
+ // stored to two decimal points.
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ int64_t extended_price = retail_price * quantity;
+ l_extendedprice[batch_offset] = {extended_price};
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_DISCOUNT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_DISCOUNT]) {
+ tld.generated_lineitem[LINEITEM::L_DISCOUNT] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 10);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_DISCOUNT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_discount =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_DISCOUNT]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_discount[batch_offset] = {dist(tld.rng)};
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_TAX(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_TAX]) {
+ tld.generated_lineitem[LINEITEM::L_TAX] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 8);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_TAX, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_tax = reinterpret_cast<Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_tax[batch_offset] = {dist(tld.rng)};
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RETURNFLAG(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_RETURNFLAG]) {
+ tld.generated_lineitem[LINEITEM::L_RETURNFLAG] = true;
+ RETURN_NOT_OK(L_RECEIPTDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<uint32_t> dist;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_RETURNFLAG, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_returnflag =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_receiptdate = reinterpret_cast<const uint32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ if (l_receiptdate[batch_offset] <= kCurrentDate) {
+ uint32_t r = dist(tld.rng);
+ l_returnflag[batch_offset] = (r % 2 == 1) ? 'R' : 'A';
+ } else {
+ l_returnflag[batch_offset] = 'N';
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINESTATUS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_LINESTATUS]) {
+ tld.generated_lineitem[LINEITEM::L_LINESTATUS] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_LINESTATUS, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_linestatus =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_LINESTATUS]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_shipdate =
+ reinterpret_cast<const uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ if (l_shipdate[batch_offset] > kCurrentDate)
+ l_linestatus[batch_offset] = 'O';
+ else
+ l_linestatus[batch_offset] = 'F';
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPDATE]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPDATE] = true;
+ RETURN_NOT_OK(O_ORDERDATE(thread_index));
+ const int32_t* o_orderdate = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(1, 121);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SHIPDATE, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_shipdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_shipdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMITDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_COMMITDATE]) {
+ tld.generated_lineitem[LINEITEM::L_COMMITDATE] = true;
+ const int32_t* o_orderdate = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(30, 90);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_COMMITDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_commitdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_COMMITDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_commitdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMITDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RECEIPTDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_RECEIPTDATE]) {
+ tld.generated_lineitem[LINEITEM::L_RECEIPTDATE] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(1, 30);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_RECEIPTDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_receiptdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_shipdate =
+ reinterpret_cast<const uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_receiptdate[batch_offset] = l_shipdate[batch_offset] + dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPINSTRUCT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT] = true;
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPINSTRUCT]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumInstructions - 1);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_SHIPINSTRUCT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_shipinstruct =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ const char* str = Instructions[dist(tld.rng)];
+ // Note that we don't have to memset the buffer to 0 because strncpy pads each
+ // string with 0's anyway
+ std::strncpy(l_shipinstruct + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPMODE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPMODE]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPMODE] = true;
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPMODE]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumModes - 1);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SHIPMODE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_shipmode =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_SHIPMODE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ const char* str = Modes[dist(tld.rng)];
+ std::strncpy(l_shipmode + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPMODE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_COMMENT]) {
+ tld.generated_lineitem[LINEITEM::L_COMMENT] = true;
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ // Comments are kind of sneaky: we always generate the full batch and then just
+ // bump the length
+ if (tld.lineitem[ibatch][LINEITEM::L_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.lineitem[ibatch][LINEITEM::L_COMMENT],
+ g_text.GenerateComments(batch_size_, 10, 43, tld.rng));
+ batch_offset = 0;
+ }
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ batch_offset += next_run;
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMENT].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ std::vector<Datum> orders;
+ int64_t orders_to_generate;
+ int64_t orderkey_start;
+
+ std::vector<std::vector<Datum>> lineitem;
+ std::vector<int> items_per_order;
+ int64_t lineitem_to_generate;
+ int64_t first_batch_offset;
+ std::bitset<LINEITEM::kNumCols> generated_lineitem;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex orders_output_queue_mutex_;
+ std::mutex lineitem_output_queue_mutex_;
+ std::queue<ExecBatch> orders_output_queue_;
+ std::queue<ExecBatch> lineitem_output_queue_;
+ int64_t batch_size_;
+ float scale_factor_;
+ int64_t orders_rows_to_generate_;
+ int64_t orders_rows_generated_;
+ std::vector<int> orders_cols_;
+ std::vector<int> lineitem_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<size_t> orders_batches_generated_ = {0};
+ std::atomic<size_t> lineitem_batches_generated_ = {0};
+};
+
+class SupplierGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 10000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, gen_list_));
+
+ random::pcg32_fast rng;
+ std::uniform_int_distribution<int64_t> dist(0, rows_to_generate_ - 1);
+ size_t num_special_rows = static_cast<size_t>(5 * scale_factor_);
+ std::unordered_set<int64_t> good_rows_set;
+ while (good_rows_set.size() < num_special_rows) {
+ int64_t row = dist(rng);
+ good_rows_set.insert(row);
+ }
+ std::unordered_set<int64_t> bad_rows_set;
+ while (bad_rows_set.size() < num_special_rows) {
+ int64_t bad_row;
+ do {
+ bad_row = dist(rng);
+ } while (good_rows_set.find(bad_row) != good_rows_set.end());
+ bad_rows_set.insert(bad_row);
+ }
+ good_rows_.clear();
+ bad_rows_.clear();
+ good_rows_.insert(good_rows_.end(), good_rows_set.begin(), good_rows_set.end());
+ bad_rows_.insert(bad_rows_.end(), bad_rows_set.begin(), bad_rows_set.end());
+ std::sort(good_rows_.begin(), good_rows_.end());
+ std::sort(bad_rows_.begin(), bad_rows_.end());
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) tld.rng.seed(seq);
+
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(S_SUPPKEY) \
+ F(S_NAME) \
+ F(S_ADDRESS) \
+ F(S_NATIONKEY) \
+ F(S_PHONE) \
+ F(S_ACCTBAL) \
+ F(S_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct SUPPLIER {
+ enum {
+ FOR_EACH_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) {#col, SUPPLIER::col},
+ const std::unordered_map<std::string, int> kNameMap = {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kGenerators = {FOR_EACH_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(), fixed_size_binary(25), utf8(),
+ int32(), fixed_size_binary(15), decimal(12, 2),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index) {
+ if (done_.load()) return Status::OK();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.suppkey_start = rows_generated_.fetch_add(batch_size_);
+ if (tld.suppkey_start >= rows_to_generate_) return Status::OK();
+
+ tld.to_generate = std::min(batch_size_, rows_to_generate_ - tld.suppkey_start);
+
+ tld.batch.resize(SUPPLIER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for (int col : gen_list_) RETURN_NOT_OK(kGenerators[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for (size_t i = 0; i < gen_list_.size(); i++) {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_outputted_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_outputted_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if (is_last_batch) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ return Status::OK();
+ }
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(kTypes[column], tld.to_generate, {nullptr, std::move(buff)});
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status S_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_SUPPKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_SUPPKEY));
+ int32_t* s_suppkey = reinterpret_cast<int32_t*>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ s_suppkey[irow] = static_cast<int32_t>(tld.suppkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_NAME].kind() == Datum::NONE) {
+ RETURN_NOT_OK(S_SUPPKEY(thread_index));
+ const int32_t* s_suppkey = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->data());
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NAME));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[SUPPLIER::S_NAME]);
+ char* s_name = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_NAME].array()->buffers[1]->mutable_data());
+ // Look man, I'm just following the spec ok? Section 4.2.3 as of March 1 2022
+ const char* supplier = "Supplie#r";
+ const size_t supplier_length = std::strlen(supplier);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = s_name + byte_width * irow;
+ std::strncpy(out, supplier, byte_width);
+ AppendNumberPaddedToNineDigits(out + supplier_length, s_suppkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ADDRESS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_ADDRESS].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[SUPPLIER::S_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status S_NATIONKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_NATIONKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t* s_nationkey = reinterpret_cast<int32_t*>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status S_PHONE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_PHONE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(S_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[SUPPLIER::S_PHONE]);
+ const int32_t* s_nationkey = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->data());
+ char* s_phone = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_PHONE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ GeneratePhoneNumber(s_phone + irow * byte_width, tld.rng, s_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ACCTBAL(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_ACCTBAL].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_ACCTBAL));
+ Decimal128* s_acctbal = reinterpret_cast<Decimal128*>(
+ tld.batch[SUPPLIER::S_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_acctbal[irow] = {dist(tld.rng)};
+ }
+ return Status::OK();
+ }
+
+ Status S_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[SUPPLIER::S_COMMENT],
+ g_text.GenerateComments(tld.to_generate, 25, 100, tld.rng));
+ ModifyComments(thread_index, "Recommends", good_rows_);
+ ModifyComments(thread_index, "Complaints", bad_rows_);
+ }
+ return Status::OK();
+ }
+
+ void ModifyComments(size_t thread_index, const char* review,
+ const std::vector<int64_t>& indices) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ const int32_t* offsets = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[1]->data());
+ char* str = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[2]->mutable_data());
+ const char* customer = "Customer";
+ const int32_t customer_length = static_cast<int32_t>(std::strlen(customer));
+ const int32_t review_length = static_cast<int32_t>(std::strlen(review));
+
+ auto it = std::lower_bound(indices.begin(), indices.end(), tld.suppkey_start);
+ for (; it != indices.end() && *it < tld.suppkey_start + tld.to_generate; it++) {
+ int64_t idx_in_batch = *it - tld.suppkey_start;
+ char* out = str + offsets[idx_in_batch];
+ int32_t str_length = offsets[idx_in_batch + 1] - offsets[idx_in_batch];
+ std::uniform_int_distribution<int32_t> gap_dist(
+ 0, str_length - customer_length - review_length);
+ int32_t gap = gap_dist(tld.rng);
+ int32_t total_length = customer_length + gap + review_length;
+ std::uniform_int_distribution<int32_t> start_dist(0, str_length - total_length);
+ int32_t start = start_dist(tld.rng);
+ std::memcpy(out + start, customer, customer_length);
+ std::memcpy(out + start + customer_length + gap, review, review_length);
+ }
+ }
+
+ struct ThreadLocalData {
+ random::pcg32_fast rng;
+ int64_t suppkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+ std::vector<int64_t> good_rows_;
+ std::vector<int64_t> bad_rows_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_;
+ std::atomic<int64_t> rows_generated_;
+ float scale_factor_;
+ int64_t batch_size_;
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class PartGenerator : public TpchTableGenerator {
+ public:
+ explicit PartGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetPartOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch, gen_->NextPartBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->part_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class PartSuppGenerator : public TpchTableGenerator {
+ public:
+ explicit PartSuppGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetPartSuppOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextPartSuppBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->partsupp_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class CustomerGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, gen_list_));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) tld.rng.seed(seq);
+
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(C_CUSTKEY) \
+ F(C_NAME) \
+ F(C_ADDRESS) \
+ F(C_NATIONKEY) \
+ F(C_PHONE) \
+ F(C_ACCTBAL) \
+ F(C_MKTSEGMENT) \
+ F(C_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct CUSTOMER {
+ enum {
+ FOR_EACH_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) {#col, CUSTOMER::col},
+ const std::unordered_map<std::string, int> kNameMap = {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kGenerators = {FOR_EACH_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(),
+ utf8(),
+ utf8(),
+ int32(),
+ fixed_size_binary(15),
+ decimal(12, 2),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index) {
+ if (done_.load()) return Status::OK();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.custkey_start = rows_generated_.fetch_add(batch_size_);
+ if (tld.custkey_start >= rows_to_generate_) return Status::OK();
+
+ tld.to_generate = std::min(batch_size_, rows_to_generate_ - tld.custkey_start);
+
+ tld.batch.resize(CUSTOMER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for (int col : gen_list_) RETURN_NOT_OK(kGenerators[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for (size_t i = 0; i < gen_list_.size(); i++) {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_generated_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_generated_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if (is_last_batch) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ return Status::OK();
+ }
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(kTypes[column], tld.to_generate, {nullptr, std::move(buff)});
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status C_CUSTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_CUSTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_CUSTKEY));
+ int32_t* c_custkey = reinterpret_cast<int32_t*>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ c_custkey[irow] = static_cast<int32_t>(tld.custkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_NAME].kind() == Datum::NONE) {
+ RETURN_NOT_OK(C_CUSTKEY(thread_index));
+ const int32_t* c_custkey = reinterpret_cast<const int32_t*>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->data());
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ const char* customer = "Customer#";
+ const size_t customer_length = std::strlen(customer);
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ int num_digits = GetNumDigits(c_custkey[irow]);
+ int num_chars = std::max(num_digits, 9);
+ offsets[irow + 1] =
+ static_cast<int32_t>(offsets[irow] + num_chars + customer_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[tld.to_generate]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = str + offsets[irow];
+ std::memcpy(out, customer, customer_length);
+ AppendNumberPaddedToNineDigits(out + customer_length, c_custkey[irow]);
+ }
+ ArrayData ad(utf8(), tld.to_generate,
+ {nullptr, std::move(offset_buff), std::move(str_buff)});
+ tld.batch[CUSTOMER::C_NAME] = std::move(ad);
+ }
+ return Status::OK();
+ }
+
+ Status C_ADDRESS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_ADDRESS].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[CUSTOMER::C_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status C_NATIONKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_NATIONKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t* c_nationkey = reinterpret_cast<int32_t*>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status C_PHONE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_PHONE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(C_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[CUSTOMER::C_PHONE]);
+ const int32_t* c_nationkey = reinterpret_cast<const int32_t*>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->data());
+ char* c_phone = reinterpret_cast<char*>(
+ tld.batch[CUSTOMER::C_PHONE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ GeneratePhoneNumber(c_phone + irow * byte_width, tld.rng, c_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_ACCTBAL(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_ACCTBAL].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_ACCTBAL));
+ Decimal128* c_acctbal = reinterpret_cast<Decimal128*>(
+ tld.batch[CUSTOMER::C_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_acctbal[irow] = {dist(tld.rng)};
+ }
+ return Status::OK();
+ }
+
+ Status C_MKTSEGMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_MKTSEGMENT].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_MKTSEGMENT));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[CUSTOMER::C_MKTSEGMENT]);
+ char* c_mktsegment = reinterpret_cast<char*>(
+ tld.batch[CUSTOMER::C_MKTSEGMENT].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int32_t> dist(0, kNumSegments - 1);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = c_mktsegment + irow * byte_width;
+ int str_idx = dist(tld.rng);
+ std::strncpy(out, Segments[str_idx], byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[CUSTOMER::C_COMMENT],
+ g_text.GenerateComments(tld.to_generate, 29, 116, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ random::pcg32_fast rng;
+ int64_t custkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_{0};
+ std::atomic<int64_t> rows_generated_ = {0};
+ float scale_factor_{0};
+ int64_t batch_size_{0};
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class OrdersGenerator : public TpchTableGenerator {
+ public:
+ explicit OrdersGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetOrdersOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch, gen_->NextOrdersBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->orders_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class LineitemGenerator : public TpchTableGenerator {
+ public:
+ explicit LineitemGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetLineItemOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextLineItemBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->lineitem_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class NationGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, float /*scale_factor*/,
+ int64_t /*batch_size*/) override {
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, column_indices_));
+ rng_.seed(arrow_vendored::pcg_extras::seed_seq_from<std::random_device>{});
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t /*num_threads*/, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback /*schedule_task_callback*/) override {
+ std::shared_ptr<Buffer> N_NATIONKEY_buffer =
+ Buffer::Wrap(kNationKey, sizeof(kNationKey));
+ ArrayData N_NATIONKEY_arraydata(int32(), kRowCount,
+ {nullptr, std::move(N_NATIONKEY_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> N_NAME_buffer,
+ AllocateBuffer(kRowCount * kNameByteWidth));
+ char* N_NAME = reinterpret_cast<char*>(N_NAME_buffer->mutable_data());
+ for (size_t i = 0; i < kRowCount; i++)
+ std::strncpy(N_NAME + kNameByteWidth * i, kCountryNames[i], kNameByteWidth);
+ ArrayData N_NAME_arraydata(fixed_size_binary(kNameByteWidth), kRowCount,
+ {nullptr, std::move(N_NAME_buffer)});
+
+ std::shared_ptr<Buffer> N_REGIONKEY_buffer =
+ Buffer::Wrap(kRegionKey, sizeof(kRegionKey));
+ ArrayData N_REGIONKEY_arraydata(int32(), kRowCount,
+ {nullptr, std::move(N_REGIONKEY_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(Datum N_COMMENT_datum,
+ g_text.GenerateComments(kRowCount, 31, 114, rng_));
+
+ std::vector<Datum> fields = {
+ std::move(N_NATIONKEY_arraydata), std::move(N_NAME_arraydata),
+ std::move(N_REGIONKEY_arraydata), std::move(N_COMMENT_datum)};
+
+ std::vector<Datum> result;
+ for (const int& col : column_indices_) result.push_back(fields[col]);
+ ARROW_ASSIGN_OR_RAISE(ExecBatch batch, ExecBatch::Make(std::move(result)));
+ output_callback(std::move(batch));
+ finished_callback(static_cast<int64_t>(1));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ random::pcg32_fast rng_;
+
+ static constexpr size_t kRowCount = 25;
+ static constexpr int32_t kNameByteWidth = 25;
+ const int32_t kNationKey[kRowCount] = {0, 1, 2, 3, 4, 5, 6, 7, 8,
+ 9, 10, 11, 12, 13, 14, 15, 16, 17,
+ 18, 19, 20, 21, 22, 23, 24};
+ const char* kCountryNames[kRowCount] = {
+ "ALGERIA", "ARGENTINA", "BRAZIL", "CANADA", "EGYPT",
+ "ETHIOPIA", "FRANCE", "GERMANY", "INDIA", "INDONESIA",
+ "IRAN", "IRAQ", "JAPAN", "JORDAN", "KENYA",
+ "MOROCCO", "MOZAMBIQUE", "PERU", "CHINA", "ROMANIA",
+ "SAUDI ARABIA", "VIETNAM", "RUSSIA", "UNITED KINGDOM", "UNITED STATES"};
+ const int32_t kRegionKey[kRowCount] = {0, 1, 1, 1, 4, 0, 3, 3, 2, 2, 4, 4, 2,
+ 4, 0, 0, 0, 1, 2, 3, 4, 2, 3, 3, 1};
+
+ struct NATION {
+ enum {
Review comment:
Also perhaps call this NationColumn or NationField?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829317934
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ if (!partsupp_output_queue_.empty()) {
+ ExecBatch result = std::move(partsupp_output_queue_.front());
+ partsupp_output_queue_.pop();
+ return std::move(result);
+ }
+ }
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+ if (!part_cols_.empty()) {
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch part_batch,
+ ExecBatch::Make(std::move(part_result)));
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ part_output_queue_.emplace(std::move(part_batch));
+ }
+ }
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ for (size_t i = 1; i < partsupp_results.size(); i++)
+ partsupp_output_queue_.emplace(std::move(partsupp_results[i]));
+ }
+ return std::move(partsupp_results[0]);
+ }
+
+ private:
+#define FOR_EACH_PART_COLUMN(F) \
+ F(P_PARTKEY) \
+ F(P_NAME) \
+ F(P_MFGR) \
+ F(P_BRAND) \
+ F(P_TYPE) \
+ F(P_SIZE) \
+ F(P_CONTAINER) \
+ F(P_RETAILPRICE) \
+ F(P_COMMENT)
+
+#define FOR_EACH_PARTSUPP_COLUMN(F) \
+ F(PS_PARTKEY) \
+ F(PS_SUPPKEY) \
+ F(PS_AVAILQTY) \
+ F(PS_SUPPLYCOST) \
+ F(PS_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct PART {
+ enum {
+ FOR_EACH_PART_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct PARTSUPP {
+ enum {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, PART::col},
+ const std::unordered_map<std::string, int> kPartNameMap = {
+ FOR_EACH_PART_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, PARTSUPP::col},
+ const std::unordered_map<std::string, int> kPartsuppNameMap = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kPartGenerators = {FOR_EACH_PART_COLUMN(MAKE_FN_ARRAY)};
+ std::vector<GenerateColumnFn> kPartsuppGenerators = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kPartTypes = {
+ int32(), utf8(), fixed_size_binary(25), fixed_size_binary(10),
+ utf8(), int32(), fixed_size_binary(10), decimal(12, 2),
+ utf8(),
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kPartsuppTypes = {
+ int32(), int32(), int32(), decimal(12, 2), utf8(),
+ };
+
+ Status AllocatePartBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.part[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.part_to_generate * byte_width));
+ ArrayData ad(kPartTypes[column], tld.part_to_generate, {nullptr, std::move(buff)});
+ tld.part[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status P_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_PARTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_PARTKEY));
+ int32_t* p_partkey = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) {
+ p_partkey[i] = static_cast<int32_t>(tld.partkey_start + i + 1);
+ ARROW_DCHECK(1 <= p_partkey[i] && p_partkey[i] <= part_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_NAME].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(
+ 0, static_cast<uint8_t>(kNumNameParts - 1));
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dist(tld.rng));
+ tld.string_indices[irow * 5 + ipart] = name_part_index;
+ string_length += std::strlen(NameParts[name_part_index]);
+ }
+ // Add 4 because there is a space between each word (i.e. four spaces)
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length + 4);
+ }
+ // Add an extra byte for the space after in the very last string.
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate] + 1));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 5 + ipart];
+ const char* part = NameParts[name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ *row++ = ' ';
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_NAME], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_NAME] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_MFGR(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_MFGR].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* manufacturer = "Manufacturer#";
+ const size_t manufacturer_length = std::strlen(manufacturer);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_MFGR));
+ char* p_mfgr = reinterpret_cast<char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ std::strncpy(p_mfgr + byte_width * irow, manufacturer, byte_width);
+ char mfgr_id = '0' + dist(tld.rng);
+ *(p_mfgr + byte_width * irow + manufacturer_length) = mfgr_id;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_BRAND(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_BRAND].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_MFGR(thread_index));
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* brand = "Brand#";
+ const size_t brand_length = std::strlen(brand);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_BRAND));
+ const char* p_mfgr = reinterpret_cast<const char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->data());
+ char* p_brand = reinterpret_cast<char*>(
+ tld.part[PART::P_BRAND].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_BRAND]);
+ int32_t mfgr_byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ const size_t mfgr_id_offset = std::strlen("Manufacturer#");
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = p_brand + byte_width * irow;
+ char mfgr_id = *(p_mfgr + irow * mfgr_byte_width + mfgr_id_offset);
+ char brand_id = '0' + dist(tld.rng);
+ std::strncpy(row, brand, byte_width);
+ *(row + brand_length) = mfgr_id;
+ *(row + brand_length + 1) = brand_id;
+ irow += 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_TYPE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_TYPE].kind() == Datum::NONE) {
+ using D = std::uniform_int_distribution<uint32_t>;
+ D dists[] = {
+ D{0, static_cast<uint8_t>(kNumTypes_1 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_2 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_3 - 1)},
+ };
+
+ const char** types[] = {Types_1, Types_2, Types_3};
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dists[ipart](tld.rng));
+ tld.string_indices[irow * 3 + ipart] = name_part_index;
+ string_length += std::strlen(types[ipart][name_part_index]);
+ }
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate]));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 3 + ipart];
+ const char* part = types[ipart][name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_TYPE], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_TYPE] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_SIZE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_SIZE].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(1, 50);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_SIZE));
+ int32_t* p_size = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_SIZE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) p_size[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status P_CONTAINER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_CONTAINER].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist1(
+ 0, static_cast<uint8_t>(kNumContainers_1 - 1));
+ std::uniform_int_distribution<int> dist2(
+ 0, static_cast<uint8_t>(kNumContainers_2 - 1));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_CONTAINER));
+ char* p_container = reinterpret_cast<char*>(
+ tld.part[PART::P_CONTAINER].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_CONTAINER]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int container1_idx = dist1(tld.rng);
+ int container2_idx = dist2(tld.rng);
+ const char* container1 = Containers_1[container1_idx];
+ const char* container2 = Containers_2[container2_idx];
+ size_t container1_length = std::strlen(container1);
+ size_t container2_length = std::strlen(container2);
+
+ char* row = p_container + byte_width * irow;
+ std::strncpy(row, container1, byte_width);
+ std::memcpy(row + container1_length, container2, container2_length);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_RETAILPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_RETAILPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_RETAILPRICE));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+ Decimal128* p_retailprice = reinterpret_cast<Decimal128*>(
+ tld.part[PART::P_RETAILPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int32_t partkey = p_partkey[irow];
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ p_retailprice[irow] = {retail_price};
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.part[PART::P_COMMENT],
+ g_text.GenerateComments(tld.part_to_generate, 5, 22, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ int64_t PartsuppBatchesToGenerate(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t num_batches = (ps_to_generate + batch_size_ - 1) / batch_size_;
+ return num_batches;
+ }
+
+ Status InitPartsupp(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.generated_partsupp.reset();
+ int64_t num_batches = PartsuppBatchesToGenerate(thread_index);
+ tld.partsupp.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.partsupp) {
+ batch.resize(PARTSUPP::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocatePartSuppBatch(size_t thread_index, size_t ibatch, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartsuppTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kPartsuppTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.partsupp[ibatch][column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status PS_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_PARTKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_PARTKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_PARTKEY));
+ int32_t* ps_partkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run; ipartsupp++, irun++)
+ ps_partkey[batch_offset++] = p_partkey[ipart];
+
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPKEY));
+ int32_t* ps_suppkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run;
+ ipartsupp++, irun++) {
+ int32_t supplier = static_cast<int32_t>(ipartsupp);
+ int32_t partkey = p_partkey[ipart];
+ ps_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
Review comment:
yes, it is magic formula.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829320809
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<float>(scale_factor));
Review comment:
changed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829327429
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
+class OrdersAndLineItemGenerator;
+class PartAndPartSupplierGenerator;
+
+class ARROW_EXPORT TpchGen {
+ public:
+ /*
+ * \brief Create a factory for nodes that generate TPC-H data
+ *
+ * Note: Individual tables will reference each other. It is important that you only
+ * create a single TpchGen instance for each plan and then you can create nodes for each
+ * table from that single TpchGen instance. Note: Every batch will be scheduled as a new
+ * task using the ExecPlan's scheduler.
+ */
+ static Result<TpchGen> Make(ExecPlan* plan, float scale_factor = 1.0f,
+ int64_t batch_size = 4096);
Review comment:
our current batch size of 32k is way way way too large; i hope to address it at some point in the near-ish future. in short, scheduling should happen on segments on the order of 1 million rows, and these segments should be broken up into batches of somewhere between 1024 and 4096 rows (ideally you want it to fit into L1 cache, which is 32 kb).
with this being at 4096, it will at least stay cache-resident for the full execution pipeline, even if scheduling overhead is a bit high.
later once we introduce separation between a batch and a unit of scheduling, we may tweak this 4096, but it's in the right ballpark for now. once that is implemented, it will be so that when we schedule the ExecNode, it will make 1M / 4096 calls to the batch producer and run the pipeline each small batch.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829354867
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
Review comment:
i will add comments. there is thread synchronization but i don't see how ThreadPool can help me avoid grabbing mutexes on shared objects.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829348174
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
Review comment:
this pseudotext is 300 MB big, and its generation is fairly expensive. i'd like to make it generated once per process.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829437850
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<float>(scale_factor));
+
+ ExecNode* lineitem =
+ *gen.Lineitem({"L_QUANTITY", "L_EXTENDEDPRICE", "L_TAX", "L_DISCOUNT", "L_SHIPDATE",
+ "L_RETURNFLAG", "L_LINESTATUS"});
+
+ std::shared_ptr<Date32Scalar> sept_2_1998 = std::make_shared<Date32Scalar>(
+ 10471); // September 2, 1998 is 10471 days after January 1, 1970
+ Expression filter =
+ less_equal(field_ref("L_SHIPDATE"), literal(std::move(sept_2_1998)));
+ FilterNodeOptions filter_opts(filter);
+
+ Expression l_returnflag = field_ref("L_RETURNFLAG");
+ Expression l_linestatus = field_ref("L_LINESTATUS");
+ Expression quantity = field_ref("L_QUANTITY");
+ Expression base_price = field_ref("L_EXTENDEDPRICE");
+
+ std::shared_ptr<Decimal128Scalar> decimal_1 =
+ std::make_shared<Decimal128Scalar>(Decimal128{0, 100}, decimal(12, 2));
+ Expression discount_multiplier =
+ call("subtract", {literal(decimal_1), field_ref("L_DISCOUNT")});
+ Expression tax_multiplier = call("add", {literal(decimal_1), field_ref("L_TAX")});
+ Expression disc_price =
+ call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier});
+ Expression charge =
+ call("multiply",
+ {call("cast",
+ {call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier})},
+ compute::CastOptions::Unsafe(decimal(12, 2))),
+ tax_multiplier});
+ Expression discount = field_ref("L_DISCOUNT");
+
+ std::vector<Expression> projection_list = {l_returnflag, l_linestatus, quantity,
+ base_price, disc_price, charge,
+ quantity, base_price, discount};
+ std::vector<std::string> project_names = {
+ "l_returnflag", "l_linestatus", "sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price", "avg_disc"};
+ ProjectNodeOptions project_opts(std::move(projection_list));
+
+ ScalarAggregateOptions sum_opts = ScalarAggregateOptions::Defaults();
+ CountOptions count_opts(CountOptions::CountMode::ALL);
+ std::vector<arrow::compute::internal::Aggregate> aggs = {
+ {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts},
+ {"hash_sum", &sum_opts}, {"hash_mean", &sum_opts}, {"hash_mean", &sum_opts},
+ {"hash_mean", &sum_opts}, {"hash_count", &count_opts}};
+
+ std::vector<FieldRef> cols = {2, 3, 4, 5, 6, 7, 8, 2};
Review comment:
why is it nicer?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829483098
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
Review comment:
This is a relatively recent introduction to the style guide (sadly it is not versioned) that we chose not to adopt: https://issues.apache.org/jira/browse/ARROW-14960
We really should add this fact to the C++ development guidelines. I will make a reminder to check it this evening.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829543526
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
Review comment:
ok, i'll change it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1057411772
@jonkeane They get generated at the same time if you use the same TpchGen object.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1057566533
Yeah that was what I was thinking. I just did a bit of tinkering with one gen node and a write node, but don’t have it working quite yet. I’ll push what I have when I get home and maybe it’ll be obvious what I’m doing wrong.
Tangentially: ideally, I would like each table to go to a single parquet file (with multiple row groups, and written batch by batch). Do we have the node machinery for that already? Or just machinery for writing to datasets?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] jonkeane commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
jonkeane commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r817267836
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3738 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status Init();
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ void GenerateWord(size_t &offset, const char **words, size_t num_choices);
+ void GenerateNoun(size_t &offset);
+ void GenerateVerb(size_t &offset);
+ void GenerateAdjective(size_t &offset);
+ void GenerateAdverb(size_t &offset);
+ void GeneratePreposition(size_t &offset);
+ void GenerateAuxiliary(size_t &offset);
+ void GenerateTerminator(size_t &offset);
+
+ void GenerateNounPhrase(size_t &offset);
+ void GenerateVerbPhrase(size_t &offset);
+ void GeneratePrepositionalPhrase(size_t &offset);
+
+ void GenerateSentence(size_t &offset);
+
+ std::unique_ptr<Buffer> text_;
+ random::pcg32_fast rng_;
+ static constexpr size_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ int scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_generated_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
Review comment:
_nods_ modified is ok, and it might be changed enough, but I'll get a second opinion on it as well to ensure we're good here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r824952090
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string> &columns,
+ const std::vector<std::shared_ptr<DataType>> &types,
+ const std::unordered_map<std::string, int> &name_map,
+ std::vector<int> &gen_list)
+ {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if(columns.empty())
+ {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for(auto pair : name_map)
+ {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ }
+ else
+ {
+ for(const std::string &col : columns)
+ {
+ auto entry = name_map.find(col);
+ if(entry == name_map.end())
+ return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+ }
+
+ static TpchText g_text;
+
+ Status TpchText::InitIfNeeded(random::pcg32_fast &rng)
+ {
+ if(done_.load())
+ return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
Review comment:
No, multiple threads can call initialization. Your version would only return Status::OK() once, which is not what we want. The lock is to ensure that only one thread allocates the 300MB chunk. I need a lock and not a `compare_exchange_strong` because I need everyone to block until the chunk is allocated.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r824953047
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
Review comment:
I need the same one to be visible to each thread.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r824957311
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
Review comment:
The spec defines generating comments as a "text string"
```
The term text string[min, max] represents a substring of a 300 MB string populated according to the pseudo text grammar defined in Clause 4.2.2.14. The length of the substring is a random number between min and max inclusive. The substring offset is randomly chosen.
```
It does also refer to it as "pseudotext" at some points, so I could name it like "TpchPseudotext" or something. I'll add a comment.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r824993200
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
Review comment:
```
bool Abort() {
bool expected = false;
return done_.compare_exchange_strong(expected, true);
}
```
Then in the node...
```
void StopProducing() override
{
if (generator_->Abort()) {
this->finished_.MarkFinished();
}
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] westonpace commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
westonpace commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r824991795
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
Review comment:
I like `TpchPseudotext`. You could also do `TpchTextString` if you want to match the spec.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r824961807
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3836 @@
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/future.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <random>
+#include <vector>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_set>
+
+namespace arrow
+{
+ using internal::checked_cast;
+
+ namespace compute
+ {
+ class TpchText
+ {
+ public:
+ Status InitIfNeeded(random::pcg32_fast &rng);
+ Result<Datum> GenerateComments(
+ size_t num_comments,
+ size_t min_length,
+ size_t max_length,
+ random::pcg32_fast &rng);
+
+ private:
+ bool GenerateWord(int64_t &offset, random::pcg32_fast &rng, char *arr, const char **words, size_t num_choices);
+ bool GenerateNoun(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdjective(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAdverb(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePreposition(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateAuxiliary(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateTerminator(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateNounPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GenerateVerbPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+ bool GeneratePrepositionalPhrase(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ bool GenerateSentence(int64_t &offset, random::pcg32_fast &rng, char *arr);
+
+ std::atomic<bool> done_ = { false };
+ int64_t generated_offset_ = 0;
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+ };
+
+ class TpchTableGenerator
+ {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(
+ std::vector<std::string> columns,
+ float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(
+ size_t num_threads,
+ OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ void Abort(AbortCallback abort_callback)
+ {
+ bool expected = false;
+ if(done_.compare_exchange_strong(expected, true))
+ {
+ abort_callback();
+ }
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = { false };
+ std::atomic<int64_t> batches_outputted_ = { 0 };
+ };
+
+ int GetNumDigits(int64_t x)
+ {
+ // This if statement chain is for MAXIMUM SPEED
+ /*
+ .,
+ . _,'f----.._
+ |\ ,-'"/ | ,'
+ |,_ ,--. /
+ /,-. ,'`. (_
+ f o| o|__ "`-.
+ ,-._.,--'_ `. _.,-`
+ `"' ___.,'` j,-'
+ `-.__.,--'
+ */
+ // Source: https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if(x < 10ll) return 1;
+ if(x < 100ll) return 2;
+ if(x < 1000ll) return 3;
+ if(x < 10000ll) return 4;
+ if(x < 100000ll) return 5;
+ if(x < 1000000ll) return 6;
+ if(x < 10000000ll) return 7;
+ if(x < 100000000ll) return 8;
+ if(x < 1000000000ll) return 9;
+ if(x < 10000000000ll) return 10;
+ if(x < 100000000000ll) return 11;
+ if(x < 1000000000000ll) return 12;
+ if(x < 10000000000000ll) return 13;
+ if(x < 100000000000000ll) return 14;
+ if(x < 1000000000000000ll) return 15;
+ if(x < 10000000000000000ll) return 16;
+ if(x < 100000000000000000ll) return 17;
+ if(x < 1000000000000000000ll) return 18;
+ return -1;
+ }
+
+ void AppendNumberPaddedToNineDigits(char *out, int64_t x)
+ {
+ // We do all of this to avoid calling snprintf, which does a lot of crazy
+ // locale stuff. On Windows and MacOS this can get suuuuper slow
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while(x > 0)
+ {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+ }
+
+ Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string> &columns,
+ const std::vector<std::shared_ptr<DataType>> &types,
+ const std::unordered_map<std::string, int> &name_map,
+ std::vector<int> &gen_list)
+ {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if(columns.empty())
+ {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for(auto pair : name_map)
+ {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ }
+ else
+ {
+ for(const std::string &col : columns)
+ {
+ auto entry = name_map.find(col);
+ if(entry == name_map.end())
+ return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+ }
+
+ static TpchText g_text;
+
+ Status TpchText::InitIfNeeded(random::pcg32_fast &rng)
+ {
+ if(done_.load())
+ return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if(!text_)
+ {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char *out = reinterpret_cast<char *>(text_->mutable_data());
+ char temp_buff[kChunkSize];
Review comment:
I don't know the length of each sentence up front, so I have to generate it first in order to know how much to increment the cursor by.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r829317347
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ if (!partsupp_output_queue_.empty()) {
+ ExecBatch result = std::move(partsupp_output_queue_.front());
+ partsupp_output_queue_.pop();
+ return std::move(result);
+ }
+ }
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+ if (!part_cols_.empty()) {
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch part_batch,
+ ExecBatch::Make(std::move(part_result)));
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ part_output_queue_.emplace(std::move(part_batch));
+ }
+ }
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ for (size_t i = 1; i < partsupp_results.size(); i++)
+ partsupp_output_queue_.emplace(std::move(partsupp_results[i]));
+ }
+ return std::move(partsupp_results[0]);
+ }
+
+ private:
+#define FOR_EACH_PART_COLUMN(F) \
+ F(P_PARTKEY) \
+ F(P_NAME) \
+ F(P_MFGR) \
+ F(P_BRAND) \
+ F(P_TYPE) \
+ F(P_SIZE) \
+ F(P_CONTAINER) \
+ F(P_RETAILPRICE) \
+ F(P_COMMENT)
+
+#define FOR_EACH_PARTSUPP_COLUMN(F) \
+ F(PS_PARTKEY) \
+ F(PS_SUPPKEY) \
+ F(PS_AVAILQTY) \
+ F(PS_SUPPLYCOST) \
+ F(PS_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct PART {
+ enum {
+ FOR_EACH_PART_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct PARTSUPP {
+ enum {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, PART::col},
+ const std::unordered_map<std::string, int> kPartNameMap = {
+ FOR_EACH_PART_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, PARTSUPP::col},
+ const std::unordered_map<std::string, int> kPartsuppNameMap = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kPartGenerators = {FOR_EACH_PART_COLUMN(MAKE_FN_ARRAY)};
+ std::vector<GenerateColumnFn> kPartsuppGenerators = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kPartTypes = {
+ int32(), utf8(), fixed_size_binary(25), fixed_size_binary(10),
+ utf8(), int32(), fixed_size_binary(10), decimal(12, 2),
+ utf8(),
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kPartsuppTypes = {
+ int32(), int32(), int32(), decimal(12, 2), utf8(),
+ };
+
+ Status AllocatePartBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.part[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.part_to_generate * byte_width));
+ ArrayData ad(kPartTypes[column], tld.part_to_generate, {nullptr, std::move(buff)});
+ tld.part[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status P_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_PARTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_PARTKEY));
+ int32_t* p_partkey = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) {
+ p_partkey[i] = static_cast<int32_t>(tld.partkey_start + i + 1);
+ ARROW_DCHECK(1 <= p_partkey[i] && p_partkey[i] <= part_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_NAME].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(
+ 0, static_cast<uint8_t>(kNumNameParts - 1));
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dist(tld.rng));
+ tld.string_indices[irow * 5 + ipart] = name_part_index;
+ string_length += std::strlen(NameParts[name_part_index]);
+ }
+ // Add 4 because there is a space between each word (i.e. four spaces)
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length + 4);
+ }
+ // Add an extra byte for the space after in the very last string.
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate] + 1));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 5 + ipart];
+ const char* part = NameParts[name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ *row++ = ' ';
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_NAME], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_NAME] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_MFGR(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_MFGR].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* manufacturer = "Manufacturer#";
+ const size_t manufacturer_length = std::strlen(manufacturer);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_MFGR));
+ char* p_mfgr = reinterpret_cast<char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ std::strncpy(p_mfgr + byte_width * irow, manufacturer, byte_width);
+ char mfgr_id = '0' + dist(tld.rng);
+ *(p_mfgr + byte_width * irow + manufacturer_length) = mfgr_id;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_BRAND(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_BRAND].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_MFGR(thread_index));
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* brand = "Brand#";
+ const size_t brand_length = std::strlen(brand);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_BRAND));
+ const char* p_mfgr = reinterpret_cast<const char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->data());
+ char* p_brand = reinterpret_cast<char*>(
+ tld.part[PART::P_BRAND].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_BRAND]);
+ int32_t mfgr_byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ const size_t mfgr_id_offset = std::strlen("Manufacturer#");
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = p_brand + byte_width * irow;
+ char mfgr_id = *(p_mfgr + irow * mfgr_byte_width + mfgr_id_offset);
+ char brand_id = '0' + dist(tld.rng);
+ std::strncpy(row, brand, byte_width);
+ *(row + brand_length) = mfgr_id;
+ *(row + brand_length + 1) = brand_id;
+ irow += 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_TYPE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_TYPE].kind() == Datum::NONE) {
+ using D = std::uniform_int_distribution<uint32_t>;
+ D dists[] = {
+ D{0, static_cast<uint8_t>(kNumTypes_1 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_2 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_3 - 1)},
+ };
+
+ const char** types[] = {Types_1, Types_2, Types_3};
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dists[ipart](tld.rng));
+ tld.string_indices[irow * 3 + ipart] = name_part_index;
+ string_length += std::strlen(types[ipart][name_part_index]);
+ }
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate]));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 3 + ipart];
+ const char* part = types[ipart][name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_TYPE], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_TYPE] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_SIZE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_SIZE].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(1, 50);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_SIZE));
+ int32_t* p_size = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_SIZE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) p_size[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status P_CONTAINER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_CONTAINER].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist1(
+ 0, static_cast<uint8_t>(kNumContainers_1 - 1));
+ std::uniform_int_distribution<int> dist2(
+ 0, static_cast<uint8_t>(kNumContainers_2 - 1));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_CONTAINER));
+ char* p_container = reinterpret_cast<char*>(
+ tld.part[PART::P_CONTAINER].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_CONTAINER]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int container1_idx = dist1(tld.rng);
+ int container2_idx = dist2(tld.rng);
+ const char* container1 = Containers_1[container1_idx];
+ const char* container2 = Containers_2[container2_idx];
+ size_t container1_length = std::strlen(container1);
+ size_t container2_length = std::strlen(container2);
+
+ char* row = p_container + byte_width * irow;
+ std::strncpy(row, container1, byte_width);
+ std::memcpy(row + container1_length, container2, container2_length);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_RETAILPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_RETAILPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_RETAILPRICE));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+ Decimal128* p_retailprice = reinterpret_cast<Decimal128*>(
+ tld.part[PART::P_RETAILPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int32_t partkey = p_partkey[irow];
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ p_retailprice[irow] = {retail_price};
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.part[PART::P_COMMENT],
+ g_text.GenerateComments(tld.part_to_generate, 5, 22, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ int64_t PartsuppBatchesToGenerate(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t num_batches = (ps_to_generate + batch_size_ - 1) / batch_size_;
+ return num_batches;
+ }
+
+ Status InitPartsupp(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.generated_partsupp.reset();
+ int64_t num_batches = PartsuppBatchesToGenerate(thread_index);
+ tld.partsupp.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.partsupp) {
+ batch.resize(PARTSUPP::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocatePartSuppBatch(size_t thread_index, size_t ibatch, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartsuppTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kPartsuppTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.partsupp[ibatch][column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status PS_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_PARTKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_PARTKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_PARTKEY));
+ int32_t* ps_partkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run; ipartsupp++, irun++)
+ ps_partkey[batch_offset++] = p_partkey[ipart];
+
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPKEY));
+ int32_t* ps_suppkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run;
+ ipartsupp++, irun++) {
+ int32_t supplier = static_cast<int32_t>(ipartsupp);
+ int32_t partkey = p_partkey[ipart];
+ ps_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_AVAILQTY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_AVAILQTY]) {
+ tld.generated_partsupp[PARTSUPP::PS_AVAILQTY] = true;
+ std::uniform_int_distribution<int32_t> dist(1, 9999);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_AVAILQTY));
+ int32_t* ps_availqty =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for (int64_t irun = 0; irun < next_run; irun++) ps_availqty[irun] = dist(tld.rng);
+
+ tld.partsupp[ibatch][PARTSUPP::PS_AVAILQTY].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPLYCOST(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPLYCOST] = true;
+ std::uniform_int_distribution<int64_t> dist(100, 100000);
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t ibatch = 0;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(
+ AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPLYCOST));
+ Decimal128* ps_supplycost =
+ reinterpret_cast<Decimal128*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+ for (int64_t irun = 0; irun < next_run; irun++)
+ ps_supplycost[irun] = {dist(tld.rng)};
+
+ tld.partsupp[ibatch][PARTSUPP::PS_SUPPLYCOST].array()->length = next_run;
+ irow += next_run;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PARTSUPP::PS_COMMENT].kind() == Datum::NONE) {
+ int64_t irow = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ int64_t num_rows = std::min(batch_size_, ps_to_generate - irow);
+ ARROW_ASSIGN_OR_RAISE(tld.partsupp[ibatch][PARTSUPP::PS_COMMENT],
+ g_text.GenerateComments(num_rows, 49, 198, tld.rng));
+ irow += num_rows;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ std::vector<Datum> part;
+ std::vector<int8_t> string_indices;
+ int64_t part_to_generate{0};
+ int64_t partkey_start{0};
+
+ std::vector<std::vector<Datum>> partsupp;
+ std::bitset<PARTSUPP::kNumCols> generated_partsupp;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex part_output_queue_mutex_;
+ std::mutex partsupp_output_queue_mutex_;
+ std::queue<ExecBatch> part_output_queue_;
+ std::queue<ExecBatch> partsupp_output_queue_;
+ int64_t batch_size_{0};
+ float scale_factor_{0};
+ int64_t part_rows_to_generate_{0};
+ int64_t part_rows_generated_{0};
+ std::vector<int> part_cols_;
+ std::vector<int> partsupp_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<int64_t> part_batches_generated_ = {0};
+ std::atomic<int64_t> partsupp_batches_generated_ = {0};
+ static constexpr int64_t kPartSuppRowsPerPart = 4;
+};
+
+class OrdersAndLineItemGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ tld.items_per_order.resize(batch_size_);
+ tld.rng.seed(seq);
+ }
+ orders_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000 * 10);
+ }
+ return Status::OK();
+ }
+
+ int64_t orders_batches_generated() const { return orders_batches_generated_.load(); }
+
+ int64_t lineitem_batches_generated() const {
+ return lineitem_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetOrdersOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kOrdersTypes, kOrdersNameMap, orders_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetLineItemOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kLineitemTypes, kLineitemNameMap, lineitem_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextOrdersBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if (!orders_output_queue_.empty()) {
+ ExecBatch batch = std::move(orders_output_queue_.front());
+ orders_output_queue_.pop();
+ return std::move(batch);
+ } else if (orders_rows_generated_ == orders_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate =
+ std::min(batch_size_, orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.first_batch_offset = 0;
+ tld.generated_lineitem.reset();
+
+ for (int col : orders_cols_) RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for (int col : lineitem_cols_) RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for (size_t i = 0; i < orders_cols_.size(); i++) {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ if (!lineitem_cols_.empty()) {
+ std::vector<ExecBatch> lineitem_results;
+ for (size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++) {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for (size_t icol = 0; icol < lineitem_cols_.size(); icol++) {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(lineitem_output_queue_mutex_);
+ for (ExecBatch& eb : lineitem_results) {
+ lineitem_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(orders_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextLineItemBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ExecBatch queued;
+ bool from_queue = false;
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ if (!lineitem_output_queue_.empty()) {
+ queued = std::move(lineitem_output_queue_.front());
+ lineitem_output_queue_.pop();
+ from_queue = true;
+ }
+ }
+ tld.first_batch_offset = 0;
+ if (from_queue) {
+ ARROW_DCHECK(queued.length <= batch_size_);
+ tld.first_batch_offset = queued.length;
+ if (queued.length == batch_size_) return std::move(queued);
+ }
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ if (orders_rows_generated_ == orders_rows_to_generate_) {
+ if (from_queue) return std::move(queued);
+ return util::nullopt;
+ }
+
+ tld.orderkey_start = orders_rows_generated_;
+ tld.orders_to_generate =
+ std::min(batch_size_, orders_rows_to_generate_ - orders_rows_generated_);
+ orders_rows_generated_ += tld.orders_to_generate;
+ orders_batches_generated_.fetch_add(1ll);
+ ARROW_DCHECK(orders_rows_generated_ <= orders_rows_to_generate_);
+ }
+ tld.orders.resize(ORDERS::kNumCols);
+ std::fill(tld.orders.begin(), tld.orders.end(), Datum());
+ RETURN_NOT_OK(GenerateRowCounts(thread_index));
+ tld.generated_lineitem.reset();
+ if (from_queue) {
+ lineitem_batches_generated_.fetch_sub(1);
+ for (size_t i = 0; i < lineitem_cols_.size(); i++)
+ if (tld.lineitem[0][lineitem_cols_[i]].kind() == Datum::NONE)
+ tld.lineitem[0][lineitem_cols_[i]] = std::move(queued[i]);
+ }
+
+ for (int col : orders_cols_) RETURN_NOT_OK(kOrdersGenerators[col](thread_index));
+ for (int col : lineitem_cols_) RETURN_NOT_OK(kLineitemGenerators[col](thread_index));
+
+ if (!orders_cols_.empty()) {
+ std::vector<Datum> orders_result(orders_cols_.size());
+ for (size_t i = 0; i < orders_cols_.size(); i++) {
+ int col_idx = orders_cols_[i];
+ orders_result[i] = tld.orders[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch orders_batch,
+ ExecBatch::Make(std::move(orders_result)));
+ {
+ std::lock_guard<std::mutex> lock(orders_output_queue_mutex_);
+ orders_output_queue_.emplace(std::move(orders_batch));
+ }
+ }
+ std::vector<ExecBatch> lineitem_results;
+ for (size_t ibatch = 0; ibatch < tld.lineitem.size(); ibatch++) {
+ std::vector<Datum> lineitem_result(lineitem_cols_.size());
+ for (size_t icol = 0; icol < lineitem_cols_.size(); icol++) {
+ int col_idx = lineitem_cols_[icol];
+ lineitem_result[icol] = tld.lineitem[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(lineitem_result)));
+ lineitem_results.emplace_back(std::move(eb));
+ }
+ lineitem_batches_generated_.fetch_add(static_cast<int64_t>(lineitem_results.size()));
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(lineitem_output_queue_mutex_);
+ for (size_t i = 1; i < lineitem_results.size(); i++)
+ lineitem_output_queue_.emplace(std::move(lineitem_results[i]));
+ }
+ return std::move(lineitem_results[0]);
+ }
+
+ private:
+#define FOR_EACH_ORDERS_COLUMN(F) \
+ F(O_ORDERKEY) \
+ F(O_CUSTKEY) \
+ F(O_ORDERSTATUS) \
+ F(O_TOTALPRICE) \
+ F(O_ORDERDATE) \
+ F(O_ORDERPRIORITY) \
+ F(O_CLERK) \
+ F(O_SHIPPRIORITY) \
+ F(O_COMMENT)
+
+#define FOR_EACH_LINEITEM_COLUMN(F) \
+ F(L_ORDERKEY) \
+ F(L_PARTKEY) \
+ F(L_SUPPKEY) \
+ F(L_LINENUMBER) \
+ F(L_QUANTITY) \
+ F(L_EXTENDEDPRICE) \
+ F(L_DISCOUNT) \
+ F(L_TAX) \
+ F(L_RETURNFLAG) \
+ F(L_LINESTATUS) \
+ F(L_SHIPDATE) \
+ F(L_COMMITDATE) \
+ F(L_RECEIPTDATE) \
+ F(L_SHIPINSTRUCT) \
+ F(L_SHIPMODE) \
+ F(L_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct ORDERS {
+ enum {
+ FOR_EACH_ORDERS_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct LINEITEM {
+ enum {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, ORDERS::col},
+ const std::unordered_map<std::string, int> kOrdersNameMap = {
+ FOR_EACH_ORDERS_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, LINEITEM::col},
+ const std::unordered_map<std::string, int> kLineitemNameMap = {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ const std::vector<GenerateColumnFn> kOrdersGenerators = {
+ FOR_EACH_ORDERS_COLUMN(MAKE_FN_ARRAY)};
+ const std::vector<GenerateColumnFn> kLineitemGenerators = {
+ FOR_EACH_LINEITEM_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kOrdersTypes = {int32(),
+ int32(),
+ fixed_size_binary(1),
+ decimal(12, 2),
+ date32(),
+ fixed_size_binary(15),
+ fixed_size_binary(15),
+ int32(),
+ utf8()};
+
+ const std::vector<std::shared_ptr<DataType>> kLineitemTypes = {
+ int32(),
+ int32(),
+ int32(),
+ int32(),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ decimal(12, 2),
+ fixed_size_binary(1),
+ fixed_size_binary(1),
+ date32(),
+ date32(),
+ date32(),
+ fixed_size_binary(25),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status AllocateOrdersBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.orders[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.orders_to_generate * byte_width));
+ ArrayData ad(kOrdersTypes[column], tld.orders_to_generate,
+ {nullptr, std::move(buff)});
+ tld.orders[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status O_ORDERKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERKEY));
+ int32_t* o_orderkey = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ int32_t orderkey_index = static_cast<int32_t>(tld.orderkey_start + i);
+ int32_t index_of_run = orderkey_index / 8;
+ int32_t index_in_run = orderkey_index % 8;
+ o_orderkey[i] = (index_of_run * 32 + index_in_run + 1);
+ ARROW_DCHECK(1 <= o_orderkey[i] && o_orderkey[i] <= 4 * orders_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CUSTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_CUSTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CUSTKEY));
+
+ // Spec says it must be a random number between 1 and SF*150000 that is not
+ // divisible by 3. Rather than repeatedly generating numbers until we get to
+ // a non-divisible-by-3 number, we just generate a number between
+ // 0 and SF * 50000 - 1, multiply by 3, and then add either 1 or 2.
+ int32_t sf_50k = static_cast<int32_t>(scale_factor_ * 50000);
+ std::uniform_int_distribution<int32_t> base_dist(0, sf_50k - 1);
+ std::uniform_int_distribution<int32_t> offset_dist(1, 2);
+ int32_t* o_custkey = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_CUSTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++)
+ o_custkey[i] = 3 * base_dist(tld.rng) + offset_dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERSTATUS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERSTATUS].kind() == Datum::NONE) {
+ RETURN_NOT_OK(L_LINESTATUS(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERSTATUS));
+
+ char* o_orderstatus = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_ORDERSTATUS].array()->buffers[1]->mutable_data());
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ bool all_f = true;
+ bool all_o = true;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ const char* l_linestatus = reinterpret_cast<const char*>(
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->buffers[1]->data());
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++, batch_offset++) {
+ all_f &= l_linestatus[batch_offset] == 'F';
+ all_o &= l_linestatus[batch_offset] == 'O';
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ ARROW_DCHECK(!(all_f && all_o));
+ if (all_f)
+ o_orderstatus[iorder] = 'F';
+ else if (all_o)
+ o_orderstatus[iorder] = 'O';
+ else
+ o_orderstatus[iorder] = 'P';
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_TOTALPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_TOTALPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(L_EXTENDEDPRICE(thread_index));
+ RETURN_NOT_OK(L_TAX(thread_index));
+ RETURN_NOT_OK(L_DISCOUNT(thread_index));
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_TOTALPRICE));
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ int64_t sum = 0;
+ Decimal128* o_totalprice = reinterpret_cast<Decimal128*>(
+ tld.orders[ORDERS::O_TOTALPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ const Decimal128* l_extendedprice = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->buffers[1]->data());
+ const Decimal128* l_tax = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->data());
+ const Decimal128* l_discount = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->buffers[1]->data());
+
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++, batch_offset++) {
+ int64_t eprice = static_cast<int64_t>(l_extendedprice[batch_offset]);
+ int64_t tax = static_cast<int64_t>(l_tax[batch_offset]);
+ int64_t discount = static_cast<int64_t>(l_discount[batch_offset]);
+ sum += (eprice * (100 + tax) * (100 - discount));
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ sum /= 100 * 100;
+ o_totalprice[iorder] = {sum};
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ batch_offset = 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERDATE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERDATE));
+
+ std::uniform_int_distribution<uint32_t> dist(kStartDate, kEndDate - 151);
+ uint32_t* o_orderdate = reinterpret_cast<uint32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) o_orderdate[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status O_ORDERPRIORITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_ORDERPRIORITY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_ORDERPRIORITY));
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_ORDERPRIORITY]);
+ std::uniform_int_distribution<int32_t> dist(0, kNumPriorities - 1);
+ char* o_orderpriority = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_ORDERPRIORITY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ const char* str = Priorities[dist(tld.rng)];
+ std::strncpy(o_orderpriority + i * byte_width, str, byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_CLERK(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_CLERK].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_CLERK));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kOrdersTypes[ORDERS::O_CLERK]);
+ int64_t max_clerk_id = static_cast<int64_t>(scale_factor_ * 1000);
+ std::uniform_int_distribution<int64_t> dist(1, max_clerk_id);
+ char* o_clerk = reinterpret_cast<char*>(
+ tld.orders[ORDERS::O_CLERK].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ const char* clerk = "Clerk#";
+ const size_t clerk_length = std::strlen(clerk);
+ int64_t clerk_number = dist(tld.rng);
+ char* output = o_clerk + i * byte_width;
+ std::strncpy(output, clerk, byte_width);
+ AppendNumberPaddedToNineDigits(output + clerk_length, clerk_number);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status O_SHIPPRIORITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_SHIPPRIORITY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateOrdersBatch(thread_index, ORDERS::O_SHIPPRIORITY));
+ int32_t* o_shippriority = reinterpret_cast<int32_t*>(
+ tld.orders[ORDERS::O_SHIPPRIORITY].array()->buffers[1]->mutable_data());
+ std::memset(o_shippriority, 0, tld.orders_to_generate * sizeof(int32_t));
+ }
+ return Status::OK();
+ }
+
+ Status O_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.orders[ORDERS::O_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.orders[ORDERS::O_COMMENT],
+ g_text.GenerateComments(tld.orders_to_generate, 19, 78, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ Status GenerateRowCounts(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ std::uniform_int_distribution<int> length_dist(1, 7);
+ tld.lineitem_to_generate = 0;
+ tld.items_per_order.clear();
+ for (int64_t i = 0; i < tld.orders_to_generate; i++) {
+ int length = length_dist(tld.rng);
+ tld.items_per_order.push_back(length);
+ tld.lineitem_to_generate += length;
+ }
+ int64_t num_batches =
+ (tld.first_batch_offset + tld.lineitem_to_generate + batch_size_ - 1) /
+ batch_size_;
+ tld.lineitem.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.lineitem) {
+ batch.resize(LINEITEM::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocateLineItemBufferIfNeeded(size_t thread_index, size_t ibatch, int column,
+ size_t& out_batch_offset) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.lineitem[ibatch][column].kind() == Datum::NONE) {
+ int32_t byte_width = arrow::internal::GetByteWidth(*kLineitemTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kLineitemTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.lineitem[ibatch][column] = std::move(ad);
+ out_batch_offset = 0;
+ }
+ if (ibatch == 0) out_batch_offset = tld.first_batch_offset;
+
+ return Status::OK();
+ }
+
+ Status L_ORDERKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_ORDERKEY]) {
+ tld.generated_lineitem[LINEITEM::L_ORDERKEY] = true;
+ RETURN_NOT_OK(O_ORDERKEY(thread_index));
+ const int32_t* o_orderkey = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_ORDERKEY, batch_offset));
+ int32_t* l_linenumber =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_ORDERKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_linenumber[batch_offset++] = o_orderkey[iorder];
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_ORDERKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_PARTKEY]) {
+ tld.generated_lineitem[LINEITEM::L_PARTKEY] = true;
+
+ size_t ibatch = 0;
+ int32_t max_partkey = static_cast<int32_t>(scale_factor_ * 200000);
+ std::uniform_int_distribution<int32_t> dist(1, max_partkey);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_PARTKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_partkey =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_partkey[batch_offset] = dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SUPPKEY]) {
+ tld.generated_lineitem[LINEITEM::L_SUPPKEY] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 3);
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset = 0;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SUPPKEY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_suppkey =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const int32_t* l_partkey = reinterpret_cast<const int32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ for (int64_t i = 0; i < next_run; i++) {
+ int32_t supplier = dist(tld.rng);
+ int32_t partkey = l_partkey[batch_offset];
+ // Fun fact: the parentheses for this expression are unbalanced in the TPC-H
+ // spec.
+ l_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SUPPKEY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINENUMBER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_LINENUMBER]) {
+ tld.generated_lineitem[LINEITEM::L_LINENUMBER] = true;
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_LINENUMBER, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ int32_t* l_linenumber =
+ reinterpret_cast<int32_t*>(tld.lineitem[ibatch][LINEITEM::L_LINENUMBER]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run;
+ iline++, irun++) {
+ l_linenumber[batch_offset++] = (iline + 1);
+ ARROW_DCHECK(1 <= (iline + 1) && (iline + 1) <= 7);
+ }
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINENUMBER].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_QUANTITY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_QUANTITY]) {
+ tld.generated_lineitem[LINEITEM::L_QUANTITY] = true;
+
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int64_t> dist(1, 50);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_QUANTITY, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_quantity =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_QUANTITY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t i = 0; i < next_run; i++) {
+ // Multiply by 100 because the type is decimal(12, 2), so the decimal goes after
+ // two digits
+ int64_t quantity = dist(tld.rng) * 100;
+ l_quantity[batch_offset++] = {quantity};
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_EXTENDEDPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE]) {
+ tld.generated_lineitem[LINEITEM::L_EXTENDEDPRICE] = true;
+ RETURN_NOT_OK(L_PARTKEY(thread_index));
+ RETURN_NOT_OK(L_QUANTITY(thread_index));
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_EXTENDEDPRICE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ const int32_t* l_partkey = reinterpret_cast<const int32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_PARTKEY].array()->buffers[1]->data());
+ const Decimal128* l_quantity = reinterpret_cast<const Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_QUANTITY].array()->buffers[1]->data());
+ Decimal128* l_extendedprice =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ int64_t partkey = static_cast<int64_t>(l_partkey[batch_offset]);
+ // Divide by 100 to recover the integer representation (not Decimal).
+ int64_t quantity = static_cast<int64_t>(l_quantity[batch_offset]) / 100;
+
+ // Spec says to divide by 100, but that happens automatically due to this being
+ // stored to two decimal points.
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ int64_t extended_price = retail_price * quantity;
+ l_extendedprice[batch_offset] = {extended_price};
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_EXTENDEDPRICE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_DISCOUNT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_DISCOUNT]) {
+ tld.generated_lineitem[LINEITEM::L_DISCOUNT] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 10);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_DISCOUNT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_discount =
+ reinterpret_cast<Decimal128*>(tld.lineitem[ibatch][LINEITEM::L_DISCOUNT]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_discount[batch_offset] = {dist(tld.rng)};
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_DISCOUNT].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_TAX(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_TAX]) {
+ tld.generated_lineitem[LINEITEM::L_TAX] = true;
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(0, 8);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_TAX, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ Decimal128* l_tax = reinterpret_cast<Decimal128*>(
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->buffers[1]->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_tax[batch_offset] = {dist(tld.rng)};
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_TAX].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RETURNFLAG(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_RETURNFLAG]) {
+ tld.generated_lineitem[LINEITEM::L_RETURNFLAG] = true;
+ RETURN_NOT_OK(L_RECEIPTDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<uint32_t> dist;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_RETURNFLAG, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_returnflag =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_receiptdate = reinterpret_cast<const uint32_t*>(
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ if (l_receiptdate[batch_offset] <= kCurrentDate) {
+ uint32_t r = dist(tld.rng);
+ l_returnflag[batch_offset] = (r % 2 == 1) ? 'R' : 'A';
+ } else {
+ l_returnflag[batch_offset] = 'N';
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RETURNFLAG].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_LINESTATUS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_LINESTATUS]) {
+ tld.generated_lineitem[LINEITEM::L_LINESTATUS] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_LINESTATUS, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_linestatus =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_LINESTATUS]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_shipdate =
+ reinterpret_cast<const uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ if (l_shipdate[batch_offset] > kCurrentDate)
+ l_linestatus[batch_offset] = 'O';
+ else
+ l_linestatus[batch_offset] = 'F';
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_LINESTATUS].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPDATE]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPDATE] = true;
+ RETURN_NOT_OK(O_ORDERDATE(thread_index));
+ const int32_t* o_orderdate = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(1, 121);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SHIPDATE, batch_offset));
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_shipdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_shipdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMITDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_COMMITDATE]) {
+ tld.generated_lineitem[LINEITEM::L_COMMITDATE] = true;
+ const int32_t* o_orderdate = reinterpret_cast<const int32_t*>(
+ tld.orders[ORDERS::O_ORDERDATE].array()->buffers[1]->data());
+ std::uniform_int_distribution<uint32_t> dist(30, 90);
+ size_t ibatch = 0;
+ size_t iorder = 0;
+ int32_t iline = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_COMMITDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_commitdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_COMMITDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; iline < tld.items_per_order[iorder] && irun < next_run; iline++, irun++)
+ l_commitdate[batch_offset++] = o_orderdate[iorder] + dist(tld.rng);
+ if (iline == tld.items_per_order[iorder]) {
+ iline = 0;
+ iorder++;
+ }
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMITDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_RECEIPTDATE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_RECEIPTDATE]) {
+ tld.generated_lineitem[LINEITEM::L_RECEIPTDATE] = true;
+ RETURN_NOT_OK(L_SHIPDATE(thread_index));
+ size_t ibatch = 0;
+ std::uniform_int_distribution<int32_t> dist(1, 30);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_RECEIPTDATE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ uint32_t* l_receiptdate =
+ reinterpret_cast<uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ const uint32_t* l_shipdate =
+ reinterpret_cast<const uint32_t*>(tld.lineitem[ibatch][LINEITEM::L_SHIPDATE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++)
+ l_receiptdate[batch_offset] = l_shipdate[batch_offset] + dist(tld.rng);
+
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_RECEIPTDATE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPINSTRUCT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPINSTRUCT] = true;
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPINSTRUCT]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumInstructions - 1);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(
+ thread_index, ibatch, LINEITEM::L_SHIPINSTRUCT, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_shipinstruct =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ const char* str = Instructions[dist(tld.rng)];
+ // Note that we don't have to memset the buffer to 0 because strncpy pads each
+ // string with 0's anyway
+ std::strncpy(l_shipinstruct + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPINSTRUCT].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_SHIPMODE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_SHIPMODE]) {
+ tld.generated_lineitem[LINEITEM::L_SHIPMODE] = true;
+ int32_t byte_width =
+ arrow::internal::GetByteWidth(*kLineitemTypes[LINEITEM::L_SHIPMODE]);
+ size_t ibatch = 0;
+ std::uniform_int_distribution<size_t> dist(0, kNumModes - 1);
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ size_t batch_offset;
+ RETURN_NOT_OK(AllocateLineItemBufferIfNeeded(thread_index, ibatch,
+ LINEITEM::L_SHIPMODE, batch_offset));
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+ char* l_shipmode =
+ reinterpret_cast<char*>(tld.lineitem[ibatch][LINEITEM::L_SHIPMODE]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+
+ for (int64_t i = 0; i < next_run; i++, batch_offset++) {
+ const char* str = Modes[dist(tld.rng)];
+ std::strncpy(l_shipmode + batch_offset * byte_width, str, byte_width);
+ }
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_SHIPMODE].array()->length =
+ static_cast<int64_t>(batch_offset);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status L_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_lineitem[LINEITEM::L_COMMENT]) {
+ tld.generated_lineitem[LINEITEM::L_COMMENT] = true;
+
+ size_t batch_offset = tld.first_batch_offset;
+ size_t ibatch = 0;
+ for (int64_t irow = 0; irow < tld.lineitem_to_generate; ibatch++) {
+ // Comments are kind of sneaky: we always generate the full batch and then just
+ // bump the length
+ if (tld.lineitem[ibatch][LINEITEM::L_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.lineitem[ibatch][LINEITEM::L_COMMENT],
+ g_text.GenerateComments(batch_size_, 10, 43, tld.rng));
+ batch_offset = 0;
+ }
+
+ int64_t remaining_in_batch = static_cast<int64_t>(batch_size_ - batch_offset);
+ int64_t next_run = std::min(tld.lineitem_to_generate - irow, remaining_in_batch);
+
+ batch_offset += next_run;
+ irow += next_run;
+ tld.lineitem[ibatch][LINEITEM::L_COMMENT].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ std::vector<Datum> orders;
+ int64_t orders_to_generate;
+ int64_t orderkey_start;
+
+ std::vector<std::vector<Datum>> lineitem;
+ std::vector<int> items_per_order;
+ int64_t lineitem_to_generate;
+ int64_t first_batch_offset;
+ std::bitset<LINEITEM::kNumCols> generated_lineitem;
+ random::pcg32_fast rng;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ bool inited_ = false;
+ std::mutex orders_output_queue_mutex_;
+ std::mutex lineitem_output_queue_mutex_;
+ std::queue<ExecBatch> orders_output_queue_;
+ std::queue<ExecBatch> lineitem_output_queue_;
+ int64_t batch_size_;
+ float scale_factor_;
+ int64_t orders_rows_to_generate_;
+ int64_t orders_rows_generated_;
+ std::vector<int> orders_cols_;
+ std::vector<int> lineitem_cols_;
+ ThreadIndexer thread_indexer_;
+
+ std::atomic<size_t> orders_batches_generated_ = {0};
+ std::atomic<size_t> lineitem_batches_generated_ = {0};
+};
+
+class SupplierGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 10000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, gen_list_));
+
+ random::pcg32_fast rng;
+ std::uniform_int_distribution<int64_t> dist(0, rows_to_generate_ - 1);
+ size_t num_special_rows = static_cast<size_t>(5 * scale_factor_);
+ std::unordered_set<int64_t> good_rows_set;
+ while (good_rows_set.size() < num_special_rows) {
+ int64_t row = dist(rng);
+ good_rows_set.insert(row);
+ }
+ std::unordered_set<int64_t> bad_rows_set;
+ while (bad_rows_set.size() < num_special_rows) {
+ int64_t bad_row;
+ do {
+ bad_row = dist(rng);
+ } while (good_rows_set.find(bad_row) != good_rows_set.end());
+ bad_rows_set.insert(bad_row);
+ }
+ good_rows_.clear();
+ bad_rows_.clear();
+ good_rows_.insert(good_rows_.end(), good_rows_set.begin(), good_rows_set.end());
+ bad_rows_.insert(bad_rows_.end(), bad_rows_set.begin(), bad_rows_set.end());
+ std::sort(good_rows_.begin(), good_rows_.end());
+ std::sort(bad_rows_.begin(), bad_rows_.end());
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) tld.rng.seed(seq);
+
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(S_SUPPKEY) \
+ F(S_NAME) \
+ F(S_ADDRESS) \
+ F(S_NATIONKEY) \
+ F(S_PHONE) \
+ F(S_ACCTBAL) \
+ F(S_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct SUPPLIER {
+ enum {
+ FOR_EACH_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) {#col, SUPPLIER::col},
+ const std::unordered_map<std::string, int> kNameMap = {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kGenerators = {FOR_EACH_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(), fixed_size_binary(25), utf8(),
+ int32(), fixed_size_binary(15), decimal(12, 2),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index) {
+ if (done_.load()) return Status::OK();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.suppkey_start = rows_generated_.fetch_add(batch_size_);
+ if (tld.suppkey_start >= rows_to_generate_) return Status::OK();
+
+ tld.to_generate = std::min(batch_size_, rows_to_generate_ - tld.suppkey_start);
+
+ tld.batch.resize(SUPPLIER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for (int col : gen_list_) RETURN_NOT_OK(kGenerators[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for (size_t i = 0; i < gen_list_.size(); i++) {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_outputted_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_outputted_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if (is_last_batch) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ return Status::OK();
+ }
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(kTypes[column], tld.to_generate, {nullptr, std::move(buff)});
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status S_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_SUPPKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_SUPPKEY));
+ int32_t* s_suppkey = reinterpret_cast<int32_t*>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ s_suppkey[irow] = static_cast<int32_t>(tld.suppkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_NAME].kind() == Datum::NONE) {
+ RETURN_NOT_OK(S_SUPPKEY(thread_index));
+ const int32_t* s_suppkey = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_SUPPKEY].array()->buffers[1]->data());
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NAME));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[SUPPLIER::S_NAME]);
+ char* s_name = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_NAME].array()->buffers[1]->mutable_data());
+ // Look man, I'm just following the spec ok? Section 4.2.3 as of March 1 2022
+ const char* supplier = "Supplie#r";
+ const size_t supplier_length = std::strlen(supplier);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = s_name + byte_width * irow;
+ std::strncpy(out, supplier, byte_width);
+ AppendNumberPaddedToNineDigits(out + supplier_length, s_suppkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ADDRESS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_ADDRESS].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[SUPPLIER::S_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status S_NATIONKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_NATIONKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t* s_nationkey = reinterpret_cast<int32_t*>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status S_PHONE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_PHONE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(S_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[SUPPLIER::S_PHONE]);
+ const int32_t* s_nationkey = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_NATIONKEY].array()->buffers[1]->data());
+ char* s_phone = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_PHONE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ GeneratePhoneNumber(s_phone + irow * byte_width, tld.rng, s_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status S_ACCTBAL(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_ACCTBAL].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, SUPPLIER::S_ACCTBAL));
+ Decimal128* s_acctbal = reinterpret_cast<Decimal128*>(
+ tld.batch[SUPPLIER::S_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ s_acctbal[irow] = {dist(tld.rng)};
+ }
+ return Status::OK();
+ }
+
+ Status S_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[SUPPLIER::S_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[SUPPLIER::S_COMMENT],
+ g_text.GenerateComments(tld.to_generate, 25, 100, tld.rng));
+ ModifyComments(thread_index, "Recommends", good_rows_);
+ ModifyComments(thread_index, "Complaints", bad_rows_);
+ }
+ return Status::OK();
+ }
+
+ void ModifyComments(size_t thread_index, const char* review,
+ const std::vector<int64_t>& indices) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ const int32_t* offsets = reinterpret_cast<const int32_t*>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[1]->data());
+ char* str = reinterpret_cast<char*>(
+ tld.batch[SUPPLIER::S_COMMENT].array()->buffers[2]->mutable_data());
+ const char* customer = "Customer";
+ const int32_t customer_length = static_cast<int32_t>(std::strlen(customer));
+ const int32_t review_length = static_cast<int32_t>(std::strlen(review));
+
+ auto it = std::lower_bound(indices.begin(), indices.end(), tld.suppkey_start);
+ for (; it != indices.end() && *it < tld.suppkey_start + tld.to_generate; it++) {
+ int64_t idx_in_batch = *it - tld.suppkey_start;
+ char* out = str + offsets[idx_in_batch];
+ int32_t str_length = offsets[idx_in_batch + 1] - offsets[idx_in_batch];
+ std::uniform_int_distribution<int32_t> gap_dist(
+ 0, str_length - customer_length - review_length);
+ int32_t gap = gap_dist(tld.rng);
+ int32_t total_length = customer_length + gap + review_length;
+ std::uniform_int_distribution<int32_t> start_dist(0, str_length - total_length);
+ int32_t start = start_dist(tld.rng);
+ std::memcpy(out + start, customer, customer_length);
+ std::memcpy(out + start + customer_length + gap, review, review_length);
+ }
+ }
+
+ struct ThreadLocalData {
+ random::pcg32_fast rng;
+ int64_t suppkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+ std::vector<int64_t> good_rows_;
+ std::vector<int64_t> bad_rows_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_;
+ std::atomic<int64_t> rows_generated_;
+ float scale_factor_;
+ int64_t batch_size_;
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class PartGenerator : public TpchTableGenerator {
+ public:
+ explicit PartGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetPartOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch, gen_->NextPartBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->part_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class PartSuppGenerator : public TpchTableGenerator {
+ public:
+ explicit PartSuppGenerator(std::shared_ptr<PartAndPartSupplierGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetPartSuppOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextPartSuppBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->partsupp_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<PartAndPartSupplierGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class CustomerGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 150000);
+ rows_generated_.store(0);
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, gen_list_));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) tld.rng.seed(seq);
+
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+#define FOR_EACH_COLUMN(F) \
+ F(C_CUSTKEY) \
+ F(C_NAME) \
+ F(C_ADDRESS) \
+ F(C_NATIONKEY) \
+ F(C_PHONE) \
+ F(C_ACCTBAL) \
+ F(C_MKTSEGMENT) \
+ F(C_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct CUSTOMER {
+ enum {
+ FOR_EACH_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+#undef MAKE_ENUM
+#define MAKE_STRING_MAP(col) {#col, CUSTOMER::col},
+ const std::unordered_map<std::string, int> kNameMap = {
+ FOR_EACH_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kGenerators = {FOR_EACH_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_COLUMN
+
+ std::vector<std::shared_ptr<DataType>> kTypes = {
+ int32(),
+ utf8(),
+ utf8(),
+ int32(),
+ fixed_size_binary(15),
+ decimal(12, 2),
+ fixed_size_binary(10),
+ utf8(),
+ };
+
+ Status ProduceCallback(size_t thread_index) {
+ if (done_.load()) return Status::OK();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.custkey_start = rows_generated_.fetch_add(batch_size_);
+ if (tld.custkey_start >= rows_to_generate_) return Status::OK();
+
+ tld.to_generate = std::min(batch_size_, rows_to_generate_ - tld.custkey_start);
+
+ tld.batch.resize(CUSTOMER::kNumCols);
+ std::fill(tld.batch.begin(), tld.batch.end(), Datum());
+ for (int col : gen_list_) RETURN_NOT_OK(kGenerators[col](thread_index));
+
+ std::vector<Datum> result(gen_list_.size());
+ for (size_t i = 0; i < gen_list_.size(); i++) {
+ int col_idx = gen_list_[i];
+ result[i] = tld.batch[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(result)));
+ int64_t batches_to_generate = (rows_to_generate_ + batch_size_ - 1) / batch_size_;
+ int64_t batches_generated_before_this_one = batches_outputted_.fetch_add(1);
+ bool is_last_batch = batches_generated_before_this_one == (batches_to_generate - 1);
+ output_callback_(std::move(eb));
+ if (is_last_batch) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ return Status::OK();
+ }
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ Status AllocateColumn(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.batch[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.to_generate * byte_width));
+ ArrayData ad(kTypes[column], tld.to_generate, {nullptr, std::move(buff)});
+ tld.batch[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status C_CUSTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_CUSTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_CUSTKEY));
+ int32_t* c_custkey = reinterpret_cast<int32_t*>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ c_custkey[irow] = static_cast<int32_t>(tld.custkey_start + irow + 1);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_NAME].kind() == Datum::NONE) {
+ RETURN_NOT_OK(C_CUSTKEY(thread_index));
+ const int32_t* c_custkey = reinterpret_cast<const int32_t*>(
+ tld.batch[CUSTOMER::C_CUSTKEY].array()->buffers[1]->data());
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ const char* customer = "Customer#";
+ const size_t customer_length = std::strlen(customer);
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ int num_digits = GetNumDigits(c_custkey[irow]);
+ int num_chars = std::max(num_digits, 9);
+ offsets[irow + 1] =
+ static_cast<int32_t>(offsets[irow] + num_chars + customer_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[tld.to_generate]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = str + offsets[irow];
+ std::memcpy(out, customer, customer_length);
+ AppendNumberPaddedToNineDigits(out + customer_length, c_custkey[irow]);
+ }
+ ArrayData ad(utf8(), tld.to_generate,
+ {nullptr, std::move(offset_buff), std::move(str_buff)});
+ tld.batch[CUSTOMER::C_NAME] = std::move(ad);
+ }
+ return Status::OK();
+ }
+
+ Status C_ADDRESS(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_ADDRESS].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[CUSTOMER::C_ADDRESS],
+ RandomVString(tld.rng, tld.to_generate, 10, 40));
+ }
+ return Status::OK();
+ }
+
+ Status C_NATIONKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_NATIONKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_NATIONKEY));
+ std::uniform_int_distribution<int32_t> dist(0, 24);
+ int32_t* c_nationkey = reinterpret_cast<int32_t*>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_nationkey[irow] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status C_PHONE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_PHONE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(C_NATIONKEY(thread_index));
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_PHONE));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[CUSTOMER::C_PHONE]);
+ const int32_t* c_nationkey = reinterpret_cast<const int32_t*>(
+ tld.batch[CUSTOMER::C_NATIONKEY].array()->buffers[1]->data());
+ char* c_phone = reinterpret_cast<char*>(
+ tld.batch[CUSTOMER::C_PHONE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ GeneratePhoneNumber(c_phone + irow * byte_width, tld.rng, c_nationkey[irow]);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_ACCTBAL(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_ACCTBAL].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_ACCTBAL));
+ Decimal128* c_acctbal = reinterpret_cast<Decimal128*>(
+ tld.batch[CUSTOMER::C_ACCTBAL].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int64_t> dist(-99999, 999999);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++)
+ c_acctbal[irow] = {dist(tld.rng)};
+ }
+ return Status::OK();
+ }
+
+ Status C_MKTSEGMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_MKTSEGMENT].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocateColumn(thread_index, CUSTOMER::C_MKTSEGMENT));
+ int32_t byte_width = arrow::internal::GetByteWidth(*kTypes[CUSTOMER::C_MKTSEGMENT]);
+ char* c_mktsegment = reinterpret_cast<char*>(
+ tld.batch[CUSTOMER::C_MKTSEGMENT].array()->buffers[1]->mutable_data());
+ std::uniform_int_distribution<int32_t> dist(0, kNumSegments - 1);
+ for (int64_t irow = 0; irow < tld.to_generate; irow++) {
+ char* out = c_mktsegment + irow * byte_width;
+ int str_idx = dist(tld.rng);
+ std::strncpy(out, Segments[str_idx], byte_width);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status C_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.batch[CUSTOMER::C_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(tld.batch[CUSTOMER::C_COMMENT],
+ g_text.GenerateComments(tld.to_generate, 29, 116, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ struct ThreadLocalData {
+ random::pcg32_fast rng;
+ int64_t custkey_start;
+ int64_t to_generate;
+ std::vector<Datum> batch;
+ };
+ std::vector<ThreadLocalData> thread_local_data_;
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t rows_to_generate_{0};
+ std::atomic<int64_t> rows_generated_ = {0};
+ float scale_factor_{0};
+ int64_t batch_size_{0};
+ std::vector<int> gen_list_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class OrdersGenerator : public TpchTableGenerator {
+ public:
+ explicit OrdersGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetOrdersOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch, gen_->NextOrdersBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->orders_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class LineitemGenerator : public TpchTableGenerator {
+ public:
+ explicit LineitemGenerator(std::shared_ptr<OrdersAndLineItemGenerator> gen)
+ : gen_(std::move(gen)) {}
+
+ Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) override {
+ scale_factor_ = scale_factor;
+ batch_size_ = batch_size;
+ ARROW_ASSIGN_OR_RAISE(schema_, gen_->SetLineItemOutputColumns(columns));
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) override {
+ RETURN_NOT_OK(gen_->Init(num_threads, batch_size_, scale_factor_));
+ output_callback_ = std::move(output_callback);
+ finished_callback_ = std::move(finished_callback);
+ schedule_callback_ = std::move(schedule_callback);
+
+ for (size_t i = 0; i < num_threads; i++)
+ RETURN_NOT_OK(schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); }));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ Status ProduceCallback(size_t) {
+ if (done_.load()) return Status::OK();
+ ARROW_ASSIGN_OR_RAISE(util::optional<ExecBatch> maybe_batch,
+ gen_->NextLineItemBatch());
+ if (!maybe_batch.has_value()) {
+ int64_t batches_generated = gen_->lineitem_batches_generated();
+ if (batches_generated == batches_outputted_.load()) {
+ bool expected = false;
+ if (done_.compare_exchange_strong(expected, true))
+ finished_callback_(batches_outputted_.load());
+ }
+ return Status::OK();
+ }
+ ExecBatch batch = std::move(*maybe_batch);
+ output_callback_(std::move(batch));
+ batches_outputted_++;
+ return schedule_callback_(
+ [this](size_t thread_index) { return this->ProduceCallback(thread_index); });
+ }
+
+ OutputBatchCallback output_callback_;
+ FinishedCallback finished_callback_;
+ ScheduleCallback schedule_callback_;
+ int64_t batch_size_;
+ float scale_factor_;
+ std::shared_ptr<OrdersAndLineItemGenerator> gen_;
+ std::shared_ptr<Schema> schema_;
+};
+
+class NationGenerator : public TpchTableGenerator {
+ public:
+ Status Init(std::vector<std::string> columns, float /*scale_factor*/,
+ int64_t /*batch_size*/) override {
+ ARROW_ASSIGN_OR_RAISE(schema_,
+ SetOutputColumns(columns, kTypes, kNameMap, column_indices_));
+ rng_.seed(arrow_vendored::pcg_extras::seed_seq_from<std::random_device>{});
+ return Status::OK();
+ }
+
+ Status StartProducing(size_t /*num_threads*/, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback /*schedule_task_callback*/) override {
+ std::shared_ptr<Buffer> N_NATIONKEY_buffer =
+ Buffer::Wrap(kNationKey, sizeof(kNationKey));
+ ArrayData N_NATIONKEY_arraydata(int32(), kRowCount,
+ {nullptr, std::move(N_NATIONKEY_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> N_NAME_buffer,
+ AllocateBuffer(kRowCount * kNameByteWidth));
+ char* N_NAME = reinterpret_cast<char*>(N_NAME_buffer->mutable_data());
+ for (size_t i = 0; i < kRowCount; i++)
+ std::strncpy(N_NAME + kNameByteWidth * i, kCountryNames[i], kNameByteWidth);
+ ArrayData N_NAME_arraydata(fixed_size_binary(kNameByteWidth), kRowCount,
+ {nullptr, std::move(N_NAME_buffer)});
+
+ std::shared_ptr<Buffer> N_REGIONKEY_buffer =
+ Buffer::Wrap(kRegionKey, sizeof(kRegionKey));
+ ArrayData N_REGIONKEY_arraydata(int32(), kRowCount,
+ {nullptr, std::move(N_REGIONKEY_buffer)});
+
+ ARROW_ASSIGN_OR_RAISE(Datum N_COMMENT_datum,
+ g_text.GenerateComments(kRowCount, 31, 114, rng_));
+
+ std::vector<Datum> fields = {
+ std::move(N_NATIONKEY_arraydata), std::move(N_NAME_arraydata),
+ std::move(N_REGIONKEY_arraydata), std::move(N_COMMENT_datum)};
+
+ std::vector<Datum> result;
+ for (const int& col : column_indices_) result.push_back(fields[col]);
+ ARROW_ASSIGN_OR_RAISE(ExecBatch batch, ExecBatch::Make(std::move(result)));
+ output_callback(std::move(batch));
+ finished_callback(static_cast<int64_t>(1));
+ return Status::OK();
+ }
+
+ std::shared_ptr<Schema> schema() const override { return schema_; }
+
+ private:
+ random::pcg32_fast rng_;
+
+ static constexpr size_t kRowCount = 25;
+ static constexpr int32_t kNameByteWidth = 25;
+ const int32_t kNationKey[kRowCount] = {0, 1, 2, 3, 4, 5, 6, 7, 8,
+ 9, 10, 11, 12, 13, 14, 15, 16, 17,
+ 18, 19, 20, 21, 22, 23, 24};
+ const char* kCountryNames[kRowCount] = {
+ "ALGERIA", "ARGENTINA", "BRAZIL", "CANADA", "EGYPT",
+ "ETHIOPIA", "FRANCE", "GERMANY", "INDIA", "INDONESIA",
+ "IRAN", "IRAQ", "JAPAN", "JORDAN", "KENYA",
+ "MOROCCO", "MOZAMBIQUE", "PERU", "CHINA", "ROMANIA",
+ "SAUDI ARABIA", "VIETNAM", "RUSSIA", "UNITED KINGDOM", "UNITED STATES"};
+ const int32_t kRegionKey[kRowCount] = {0, 1, 1, 1, 4, 0, 3, 3, 2, 2, 4, 4, 2,
+ 4, 0, 0, 0, 1, 2, 3, 4, 2, 3, 3, 1};
+
+ struct NATION {
+ enum {
Review comment:
no, i want NATION::NATIONKEY et. al to be implicitly convertible to int, so that I can use it as a name of indices. with enum class you have to add `static_cast` everywhere.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ if (!partsupp_output_queue_.empty()) {
+ ExecBatch result = std::move(partsupp_output_queue_.front());
+ partsupp_output_queue_.pop();
+ return std::move(result);
+ }
+ }
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+ if (!part_cols_.empty()) {
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch part_batch,
+ ExecBatch::Make(std::move(part_result)));
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ part_output_queue_.emplace(std::move(part_batch));
+ }
+ }
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ for (size_t i = 1; i < partsupp_results.size(); i++)
+ partsupp_output_queue_.emplace(std::move(partsupp_results[i]));
+ }
+ return std::move(partsupp_results[0]);
+ }
+
+ private:
+#define FOR_EACH_PART_COLUMN(F) \
+ F(P_PARTKEY) \
+ F(P_NAME) \
+ F(P_MFGR) \
+ F(P_BRAND) \
+ F(P_TYPE) \
+ F(P_SIZE) \
+ F(P_CONTAINER) \
+ F(P_RETAILPRICE) \
+ F(P_COMMENT)
+
+#define FOR_EACH_PARTSUPP_COLUMN(F) \
+ F(PS_PARTKEY) \
+ F(PS_SUPPKEY) \
+ F(PS_AVAILQTY) \
+ F(PS_SUPPLYCOST) \
+ F(PS_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct PART {
+ enum {
+ FOR_EACH_PART_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct PARTSUPP {
+ enum {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, PART::col},
+ const std::unordered_map<std::string, int> kPartNameMap = {
+ FOR_EACH_PART_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, PARTSUPP::col},
+ const std::unordered_map<std::string, int> kPartsuppNameMap = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kPartGenerators = {FOR_EACH_PART_COLUMN(MAKE_FN_ARRAY)};
+ std::vector<GenerateColumnFn> kPartsuppGenerators = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kPartTypes = {
+ int32(), utf8(), fixed_size_binary(25), fixed_size_binary(10),
+ utf8(), int32(), fixed_size_binary(10), decimal(12, 2),
+ utf8(),
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kPartsuppTypes = {
+ int32(), int32(), int32(), decimal(12, 2), utf8(),
+ };
+
+ Status AllocatePartBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.part[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.part_to_generate * byte_width));
+ ArrayData ad(kPartTypes[column], tld.part_to_generate, {nullptr, std::move(buff)});
+ tld.part[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status P_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_PARTKEY].kind() == Datum::NONE) {
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_PARTKEY));
+ int32_t* p_partkey = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) {
+ p_partkey[i] = static_cast<int32_t>(tld.partkey_start + i + 1);
+ ARROW_DCHECK(1 <= p_partkey[i] && p_partkey[i] <= part_rows_to_generate_);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_NAME(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_NAME].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(
+ 0, static_cast<uint8_t>(kNumNameParts - 1));
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dist(tld.rng));
+ tld.string_indices[irow * 5 + ipart] = name_part_index;
+ string_length += std::strlen(NameParts[name_part_index]);
+ }
+ // Add 4 because there is a space between each word (i.e. four spaces)
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length + 4);
+ }
+ // Add an extra byte for the space after in the very last string.
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate] + 1));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 5; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 5 + ipart];
+ const char* part = NameParts[name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ *row++ = ' ';
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_NAME], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_NAME] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_MFGR(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_MFGR].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* manufacturer = "Manufacturer#";
+ const size_t manufacturer_length = std::strlen(manufacturer);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_MFGR));
+ char* p_mfgr = reinterpret_cast<char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ std::strncpy(p_mfgr + byte_width * irow, manufacturer, byte_width);
+ char mfgr_id = '0' + dist(tld.rng);
+ *(p_mfgr + byte_width * irow + manufacturer_length) = mfgr_id;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_BRAND(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_BRAND].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_MFGR(thread_index));
+ std::uniform_int_distribution<int> dist(1, 5);
+ const char* brand = "Brand#";
+ const size_t brand_length = std::strlen(brand);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_BRAND));
+ const char* p_mfgr = reinterpret_cast<const char*>(
+ tld.part[PART::P_MFGR].array()->buffers[1]->data());
+ char* p_brand = reinterpret_cast<char*>(
+ tld.part[PART::P_BRAND].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_BRAND]);
+ int32_t mfgr_byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_MFGR]);
+ const size_t mfgr_id_offset = std::strlen("Manufacturer#");
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = p_brand + byte_width * irow;
+ char mfgr_id = *(p_mfgr + irow * mfgr_byte_width + mfgr_id_offset);
+ char brand_id = '0' + dist(tld.rng);
+ std::strncpy(row, brand, byte_width);
+ *(row + brand_length) = mfgr_id;
+ *(row + brand_length + 1) = brand_id;
+ irow += 0;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_TYPE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_TYPE].kind() == Datum::NONE) {
+ using D = std::uniform_int_distribution<uint32_t>;
+ D dists[] = {
+ D{0, static_cast<uint8_t>(kNumTypes_1 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_2 - 1)},
+ D{0, static_cast<uint8_t>(kNumTypes_3 - 1)},
+ };
+
+ const char** types[] = {Types_1, Types_2, Types_3};
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((tld.part_to_generate + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ size_t string_length = 0;
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = static_cast<uint8_t>(dists[ipart](tld.rng));
+ tld.string_indices[irow * 3 + ipart] = name_part_index;
+ string_length += std::strlen(types[ipart][name_part_index]);
+ }
+ offsets[irow + 1] = static_cast<int32_t>(offsets[irow] + string_length);
+ }
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> string_buffer,
+ AllocateBuffer(offsets[tld.part_to_generate]));
+ char* strings = reinterpret_cast<char*>(string_buffer->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ char* row = strings + offsets[irow];
+ for (int ipart = 0; ipart < 3; ipart++) {
+ uint8_t name_part_index = tld.string_indices[irow * 3 + ipart];
+ const char* part = types[ipart][name_part_index];
+ size_t length = std::strlen(part);
+ std::memcpy(row, part, length);
+ row += length;
+ }
+ }
+ ArrayData ad(kPartTypes[PART::P_TYPE], tld.part_to_generate,
+ {nullptr, std::move(offset_buff), std::move(string_buffer)});
+ Datum datum(ad);
+ tld.part[PART::P_TYPE] = std::move(datum);
+ }
+ return Status::OK();
+ }
+
+ Status P_SIZE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_SIZE].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int32_t> dist(1, 50);
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_SIZE));
+ int32_t* p_size = reinterpret_cast<int32_t*>(
+ tld.part[PART::P_SIZE].array()->buffers[1]->mutable_data());
+ for (int64_t i = 0; i < tld.part_to_generate; i++) p_size[i] = dist(tld.rng);
+ }
+ return Status::OK();
+ }
+
+ Status P_CONTAINER(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_CONTAINER].kind() == Datum::NONE) {
+ std::uniform_int_distribution<int> dist1(
+ 0, static_cast<uint8_t>(kNumContainers_1 - 1));
+ std::uniform_int_distribution<int> dist2(
+ 0, static_cast<uint8_t>(kNumContainers_2 - 1));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_CONTAINER));
+ char* p_container = reinterpret_cast<char*>(
+ tld.part[PART::P_CONTAINER].array()->buffers[1]->mutable_data());
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[PART::P_CONTAINER]);
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int container1_idx = dist1(tld.rng);
+ int container2_idx = dist2(tld.rng);
+ const char* container1 = Containers_1[container1_idx];
+ const char* container2 = Containers_2[container2_idx];
+ size_t container1_length = std::strlen(container1);
+ size_t container2_length = std::strlen(container2);
+
+ char* row = p_container + byte_width * irow;
+ std::strncpy(row, container1, byte_width);
+ std::memcpy(row + container1_length, container2, container2_length);
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_RETAILPRICE(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_RETAILPRICE].kind() == Datum::NONE) {
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ RETURN_NOT_OK(AllocatePartBatch(thread_index, PART::P_RETAILPRICE));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+ Decimal128* p_retailprice = reinterpret_cast<Decimal128*>(
+ tld.part[PART::P_RETAILPRICE].array()->buffers[1]->mutable_data());
+ for (int64_t irow = 0; irow < tld.part_to_generate; irow++) {
+ int32_t partkey = p_partkey[irow];
+ int64_t retail_price =
+ (90000 + ((partkey / 10) % 20001) + 100 * (partkey % 1000));
+ p_retailprice[irow] = {retail_price};
+ }
+ }
+ return Status::OK();
+ }
+
+ Status P_COMMENT(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (tld.part[PART::P_COMMENT].kind() == Datum::NONE) {
+ ARROW_ASSIGN_OR_RAISE(
+ tld.part[PART::P_COMMENT],
+ g_text.GenerateComments(tld.part_to_generate, 5, 22, tld.rng));
+ }
+ return Status::OK();
+ }
+
+ int64_t PartsuppBatchesToGenerate(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ int64_t num_batches = (ps_to_generate + batch_size_ - 1) / batch_size_;
+ return num_batches;
+ }
+
+ Status InitPartsupp(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ tld.generated_partsupp.reset();
+ int64_t num_batches = PartsuppBatchesToGenerate(thread_index);
+ tld.partsupp.resize(num_batches);
+ for (std::vector<Datum>& batch : tld.partsupp) {
+ batch.resize(PARTSUPP::kNumCols);
+ std::fill(batch.begin(), batch.end(), Datum());
+ }
+ return Status::OK();
+ }
+
+ Status AllocatePartSuppBatch(size_t thread_index, size_t ibatch, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartsuppTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(batch_size_ * byte_width));
+ ArrayData ad(kPartsuppTypes[column], batch_size_, {nullptr, std::move(buff)});
+ tld.partsupp[ibatch][column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status PS_PARTKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_PARTKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_PARTKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_PARTKEY));
+ int32_t* ps_partkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run; ipartsupp++, irun++)
+ ps_partkey[batch_offset++] = p_partkey[ipart];
+
+ if (ipartsupp == kPartSuppRowsPerPart) {
+ ipartsupp = 0;
+ ipart++;
+ }
+ }
+ irow += next_run;
+ tld.partsupp[ibatch][PARTSUPP::PS_PARTKEY].array()->length = batch_offset;
+ }
+ }
+ return Status::OK();
+ }
+
+ Status PS_SUPPKEY(size_t thread_index) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ if (!tld.generated_partsupp[PARTSUPP::PS_SUPPKEY]) {
+ tld.generated_partsupp[PARTSUPP::PS_SUPPKEY] = true;
+ RETURN_NOT_OK(P_PARTKEY(thread_index));
+ const int32_t* p_partkey = reinterpret_cast<const int32_t*>(
+ tld.part[PART::P_PARTKEY].array()->buffers[1]->data());
+
+ size_t ibatch = 0;
+ int64_t ipartsupp = 0;
+ int64_t ipart = 0;
+ int64_t ps_to_generate = kPartSuppRowsPerPart * tld.part_to_generate;
+ const int32_t S = static_cast<int32_t>(scale_factor_ * 10000);
+ for (int64_t irow = 0; irow < ps_to_generate; ibatch++) {
+ RETURN_NOT_OK(AllocatePartSuppBatch(thread_index, ibatch, PARTSUPP::PS_SUPPKEY));
+ int32_t* ps_suppkey =
+ reinterpret_cast<int32_t*>(tld.partsupp[ibatch][PARTSUPP::PS_SUPPKEY]
+ .array()
+ ->buffers[1]
+ ->mutable_data());
+ int64_t next_run = std::min(batch_size_, ps_to_generate - irow);
+
+ int64_t batch_offset = 0;
+ for (int64_t irun = 0; irun < next_run;) {
+ for (; ipartsupp < kPartSuppRowsPerPart && irun < next_run;
+ ipartsupp++, irun++) {
+ int32_t supplier = static_cast<int32_t>(ipartsupp);
+ int32_t partkey = p_partkey[ipart];
+ ps_suppkey[batch_offset++] =
+ (partkey + (supplier * ((S / 4) + (partkey - 1) / S))) % S + 1;
Review comment:
yes, it is magic formula.
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<float>(scale_factor));
+
+ ExecNode* lineitem =
+ *gen.Lineitem({"L_QUANTITY", "L_EXTENDEDPRICE", "L_TAX", "L_DISCOUNT", "L_SHIPDATE",
+ "L_RETURNFLAG", "L_LINESTATUS"});
+
+ std::shared_ptr<Date32Scalar> sept_2_1998 = std::make_shared<Date32Scalar>(
+ 10471); // September 2, 1998 is 10471 days after January 1, 1970
+ Expression filter =
+ less_equal(field_ref("L_SHIPDATE"), literal(std::move(sept_2_1998)));
+ FilterNodeOptions filter_opts(filter);
+
+ Expression l_returnflag = field_ref("L_RETURNFLAG");
+ Expression l_linestatus = field_ref("L_LINESTATUS");
+ Expression quantity = field_ref("L_QUANTITY");
+ Expression base_price = field_ref("L_EXTENDEDPRICE");
+
+ std::shared_ptr<Decimal128Scalar> decimal_1 =
+ std::make_shared<Decimal128Scalar>(Decimal128{0, 100}, decimal(12, 2));
+ Expression discount_multiplier =
+ call("subtract", {literal(decimal_1), field_ref("L_DISCOUNT")});
+ Expression tax_multiplier = call("add", {literal(decimal_1), field_ref("L_TAX")});
+ Expression disc_price =
+ call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier});
+ Expression charge =
+ call("multiply",
+ {call("cast",
+ {call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier})},
+ compute::CastOptions::Unsafe(decimal(12, 2))),
+ tax_multiplier});
+ Expression discount = field_ref("L_DISCOUNT");
+
+ std::vector<Expression> projection_list = {l_returnflag, l_linestatus, quantity,
+ base_price, disc_price, charge,
+ quantity, base_price, discount};
+ std::vector<std::string> project_names = {
+ "l_returnflag", "l_linestatus", "sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price", "avg_disc"};
+ ProjectNodeOptions project_opts(std::move(projection_list));
+
+ ScalarAggregateOptions sum_opts = ScalarAggregateOptions::Defaults();
+ CountOptions count_opts(CountOptions::CountMode::ALL);
+ std::vector<arrow::compute::internal::Aggregate> aggs = {
+ {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts},
+ {"hash_sum", &sum_opts}, {"hash_mean", &sum_opts}, {"hash_mean", &sum_opts},
+ {"hash_mean", &sum_opts}, {"hash_count", &count_opts}};
+
+ std::vector<FieldRef> cols = {2, 3, 4, 5, 6, 7, 8, 2};
+
+ std::vector<std::string> names = {"sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price",
+ "avg_disc", "count_order"};
+
+ std::vector<FieldRef> keys = {"L_RETURNFLAG", "L_LINESTATUS"};
+ AggregateNodeOptions agg_opts(aggs, cols, names, keys);
+
+ SortKey l_returnflag_key("L_RETURNFLAG");
+ SortKey l_linestatus_key("L_LINESTATUS");
+ SortOptions sort_opts({l_returnflag_key, l_linestatus_key});
+ OrderBySinkNodeOptions order_by_opts(sort_opts, sink_gen);
+
+ Declaration filter_decl("filter", {Declaration::Input(lineitem)}, filter_opts);
+ Declaration project_decl("project", project_opts);
+ Declaration aggregate_decl("aggregate", agg_opts);
+ Declaration orderby_decl("order_by_sink", order_by_opts);
+
+ Declaration q1 =
+ Declaration::Sequence({filter_decl, project_decl, aggregate_decl, orderby_decl});
+ std::ignore = *q1.AddToPlan(plan.get());
+ return plan;
+}
+
+static void BM_Tpch_Q1(benchmark::State& st) {
+ for (auto _ : st) {
+ st.PauseTiming();
+ AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+ std::shared_ptr<ExecPlan> plan = Plan_Q1(&sink_gen, static_cast<int>(st.range(0)));
+ st.ResumeTiming();
+ auto fut = StartAndCollect(plan.get(), sink_gen);
+ auto res = *fut.MoveResult();
+ }
+}
+
+BENCHMARK(BM_Tpch_Q1)->Args({1})->ArgNames({"SF"});
Review comment:
i think if there's a space, google benchmark will add quotes around it. I could switch it to `ScaleFactor` as a compromise.
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
Review comment:
none of the other benchmarks use it, so i just copied that. I can try it though
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<float>(scale_factor));
Review comment:
changed
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
+class OrdersAndLineItemGenerator;
+class PartAndPartSupplierGenerator;
+
+class ARROW_EXPORT TpchGen {
+ public:
+ /*
+ * \brief Create a factory for nodes that generate TPC-H data
+ *
+ * Note: Individual tables will reference each other. It is important that you only
+ * create a single TpchGen instance for each plan and then you can create nodes for each
+ * table from that single TpchGen instance. Note: Every batch will be scheduled as a new
+ * task using the ExecPlan's scheduler.
+ */
+ static Result<TpchGen> Make(ExecPlan* plan, float scale_factor = 1.0f,
+ int64_t batch_size = 4096);
Review comment:
our current batch size of 32k is way way way too large; i hope to address it at some point in the near-ish future. in short, scheduling should happen on segments on the order of 1 million rows, and these segments should be broken up into batches of somewhere between 1024 and 4096 rows (ideally you want it to fit into L1 cache, which is 32 kb).
with this being at 4096, it will at least stay cache-resident for the full execution pipeline, even if scheduling overhead is a bit high.
later once we introduce separation between a batch and a unit of scheduling, we may tweak this 4096, but it's in the right ballpark for now. once that is implemented, it will be so that when we schedule the ExecNode, it will make 1M / 4096 calls to the batch producer and run the pipeline each small batch.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
+class OrdersAndLineItemGenerator;
+class PartAndPartSupplierGenerator;
+
+class ARROW_EXPORT TpchGen {
+ public:
+ /*
+ * \brief Create a factory for nodes that generate TPC-H data
+ *
+ * Note: Individual tables will reference each other. It is important that you only
+ * create a single TpchGen instance for each plan and then you can create nodes for each
+ * table from that single TpchGen instance. Note: Every batch will be scheduled as a new
+ * task using the ExecPlan's scheduler.
+ */
+ static Result<TpchGen> Make(ExecPlan* plan, float scale_factor = 1.0f,
+ int64_t batch_size = 4096);
+
+ Result<ExecNode*> Supplier(std::vector<std::string> columns = {});
Review comment:
done
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
Review comment:
done
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
Review comment:
from google style guide:
```
Parameters are either inputs to the function, outputs from the function, or both. Non-optional input parameters should usually be values or const references, while non-optional output and input/output parameters should usually be references (which cannot be null). Generally, use std::optional to represent optional by-value inputs, and use a const pointer when the non-optional form would have used a reference. Use non-const pointers to represent optional outputs and optional input/output parameters.
```
this is a non-optional input/output parameter, so according to them it should be a reference.
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
Review comment:
from google style guide:
> Parameters are either inputs to the function, outputs from the function, or both. Non-optional input parameters should usually be values or const references, while non-optional output and input/output parameters should usually be references (which cannot be null). Generally, use std::optional to represent optional by-value inputs, and use a const pointer when the non-optional form would have used a reference. Use non-const pointers to represent optional outputs and optional input/output parameters.
this is a non-optional input/output parameter, so according to them it should be a reference.
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
Review comment:
switched it to ValidateFull.
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row + num_offset;
+ int64_t num = 0;
+ int ibyte = static_cast<int>(num_offset);
+ for (; *num_str && ibyte < byte_width; ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ if (verify_padding) {
+ int num_chars = ibyte - num_offset;
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ VerifyStringAndNumber_Single(row, prefix, i, nums, byte_width, verify_padding);
+ }
+}
+
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ char tmp_str[256] = {};
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ std::memcpy(tmp_str, str + start, str_len);
+ VerifyStringAndNumber_Single(tmp_str, prefix, i, nums, sizeof(tmp_str),
+ verify_padding);
+ }
+}
+
+void VerifyVString(const Datum& d, int min_length, int max_length) {
+ int64_t length = d.length();
+ const int32_t* off = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t start = off[i];
+ int32_t end = off[i + 1];
+ int32_t str_len = end - start;
+ ASSERT_LE(str_len, max_length);
+ ASSERT_GE(str_len, min_length);
+ for (int32_t i = start; i < end; i++) {
+ bool is_valid =
+ std::isdigit(str[i]) || std::isalpha(str[i]) || str[i] == ',' || str[i] == ' ';
+ ASSERT_TRUE(is_valid) << "Character " << str[i]
+ << " is not a digit, a letter, a comma, or a space";
+ }
+ }
+}
+
+void VerifyModuloBetween(const Datum& d, int32_t min, int32_t max, int32_t mod) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int32_t m = n[i] % mod;
+ ASSERT_GE(m, min) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ ASSERT_LE(m, max) << "Value must be between " << min << " and " << max << " mod "
+ << mod << ", " << n[i] << " % " << mod << " = " << m;
+ }
+}
+
+void VerifyAllBetween(const Datum& d, int32_t min, int32_t max) {
+ int64_t length = d.length();
+ const int32_t* n = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ ASSERT_GE(n[i], min) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ ASSERT_LE(n[i], max) << "Value must be between " << min << " and " << max << ", got "
+ << n[i];
+ }
+}
+
+void VerifyNationKey(const Datum& d) { VerifyAllBetween(d, 0, 24); }
+
+void VerifyPhone(const Datum& d) {
+ int64_t length = d.length();
+ const char* phones = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ constexpr int kByteWidth = 15; // This is common for all PHONE columns
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = phones + i * kByteWidth;
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_EQ(*row++, '-');
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ ASSERT_TRUE(std::isdigit(*row++));
+ }
+}
+
+void VerifyDecimalsBetween(const Datum& d, int64_t min, int64_t max) {
+ int64_t length = d.length();
+ const Decimal128* decs =
+ reinterpret_cast<const Decimal128*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ int64_t val = static_cast<int64_t>(decs[i]);
+ ASSERT_LE(val, max);
+ ASSERT_GE(val, min);
+ }
+}
+
+void VerifyCorrectNumberOfWords_Varlen(const Datum& d, int num_words) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+
+ int32_t start = offsets[i];
+ int32_t end = offsets[i + 1];
+ int32_t str_len = end - start;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, str + start, str_len);
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = offsets[i]; j < offsets[i + 1]; j++) {
+ bool is_space = str[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(str[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << tmp_str;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << tmp_str;
+ }
+}
+
+void VerifyCorrectNumberOfWords_FixedWidth(const Datum& d, int num_words,
+ int byte_width) {
+ int expected_num_spaces = num_words - 1;
+ int64_t length = d.length();
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+
+ for (int64_t i = 0; i < length; i++) {
+ int actual_num_spaces = 0;
+ const char* row = str + i * byte_width;
+ bool is_only_alphas_or_spaces = true;
+ for (int32_t j = 0; j < byte_width && row[j]; j++) {
+ bool is_space = row[j] == ' ';
+ actual_num_spaces += is_space;
+ is_only_alphas_or_spaces &= (is_space || std::isalpha(row[j]));
+ }
+ ASSERT_TRUE(is_only_alphas_or_spaces)
+ << "Words must be composed only of letters, got " << row;
+ ASSERT_EQ(actual_num_spaces, expected_num_spaces)
+ << "Wrong number of spaces in " << row;
+ }
+}
+
+void VerifyOneOf(const Datum& d, const std::unordered_set<char>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++)
+ ASSERT_TRUE(possibilities.find(col[i]) != possibilities.end());
+}
+
+void VerifyOneOf(const Datum& d, int32_t byte_width,
+ const std::unordered_set<std::string>& possibilities) {
+ int64_t length = d.length();
+ const char* col = reinterpret_cast<const char*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = col + i * byte_width;
+ char tmp_str[256] = {};
+ std::memcpy(tmp_str, row, byte_width);
+ ASSERT_TRUE(possibilities.find(tmp_str) != possibilities.end())
+ << tmp_str << " is not a valid string.";
+ }
+}
+
+void CountInstances(std::unordered_map<int32_t, int32_t>& counts, const Datum& d) {
+ int64_t length = d.length();
+ const int32_t* nums = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ for (int64_t i = 0; i < length; i++) counts[nums[i]]++;
+}
+
+void CountModifiedComments(const Datum& d, int& good_count, int& bad_count) {
+ int64_t length = d.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(d.array()->buffers[2]->data());
+ // Length of S_COMMENT is at most 100
+ char tmp_string[101];
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + offsets[i];
+ int32_t row_length = offsets[i + 1] - offsets[i];
+ std::memset(tmp_string, 0, sizeof(tmp_string));
+ std::memcpy(tmp_string, row, row_length);
+ char* customer = std::strstr(tmp_string, "Customer");
+ char* recommends = std::strstr(tmp_string, "Recommends");
+ char* complaints = std::strstr(tmp_string, "Complaints");
+ if (customer) {
+ ASSERT_TRUE((recommends != nullptr) ^ (complaints != nullptr));
+ if (recommends) good_count++;
+ if (complaints) bad_count++;
+ }
+ }
+}
+
+TEST(TpchNode, ScaleFactor) {
+ ExecContext ctx(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(&ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), 0.25f);
+ ExecNode* table = *gen.Supplier();
+ AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+ Declaration sink("sink", {Declaration::Input(table)}, SinkNodeOptions{&sink_gen});
+ std::ignore = *sink.AddToPlan(plan.get());
+ auto fut = StartAndCollect(plan.get(), sink_gen);
+ auto res = *fut.MoveResult();
+
+ int64_t kExpectedRows = 2500;
+ int64_t num_rows = 0;
+ for (auto& batch : res) num_rows += batch.length;
+ ASSERT_EQ(num_rows, kExpectedRows);
+ arrow::internal::GetCpuThreadPool()->WaitForIdle();
Review comment:
it was necessary at one point but i think it's not anymore. i've deleted it.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
Review comment:
this pseudotext is 300 MB big, and its generation is fairly expensive. i'd like to make it generated once per process.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
Review comment:
yes, the spec specifies to generate in the range 100-999 and 1000-9999
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
Review comment:
i will add comments. there is thread synchronization but i don't see how ThreadPool can help me avoid grabbing mutexes on shared objects.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
Review comment:
i switched it to use GetRandomSeed. in regards to nondeterminism between runs i'll let Weston decide.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
Review comment:
i don't want all of this stuff to be heap allocated if it doesn't need to be (which it doesn't). i could do `std::array<const char *>` but that is also a pain because then i have to count the number of elements manually. i think they only added the template argument deduction for the size in C++17.
##########
File path: cpp/src/arrow/compute/exec/tpch_node.h
##########
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/util/pcg_random.h"
+
+namespace arrow {
+namespace compute {
Review comment:
actually nevermind, i think users should be able to use this class for their exec plans. this can be a provided data generator.
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<float>(scale_factor));
+
+ ExecNode* lineitem =
+ *gen.Lineitem({"L_QUANTITY", "L_EXTENDEDPRICE", "L_TAX", "L_DISCOUNT", "L_SHIPDATE",
+ "L_RETURNFLAG", "L_LINESTATUS"});
+
+ std::shared_ptr<Date32Scalar> sept_2_1998 = std::make_shared<Date32Scalar>(
+ 10471); // September 2, 1998 is 10471 days after January 1, 1970
+ Expression filter =
+ less_equal(field_ref("L_SHIPDATE"), literal(std::move(sept_2_1998)));
+ FilterNodeOptions filter_opts(filter);
+
+ Expression l_returnflag = field_ref("L_RETURNFLAG");
+ Expression l_linestatus = field_ref("L_LINESTATUS");
+ Expression quantity = field_ref("L_QUANTITY");
+ Expression base_price = field_ref("L_EXTENDEDPRICE");
+
+ std::shared_ptr<Decimal128Scalar> decimal_1 =
+ std::make_shared<Decimal128Scalar>(Decimal128{0, 100}, decimal(12, 2));
+ Expression discount_multiplier =
+ call("subtract", {literal(decimal_1), field_ref("L_DISCOUNT")});
+ Expression tax_multiplier = call("add", {literal(decimal_1), field_ref("L_TAX")});
+ Expression disc_price =
+ call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier});
+ Expression charge =
+ call("multiply",
+ {call("cast",
+ {call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier})},
+ compute::CastOptions::Unsafe(decimal(12, 2))),
+ tax_multiplier});
+ Expression discount = field_ref("L_DISCOUNT");
+
+ std::vector<Expression> projection_list = {l_returnflag, l_linestatus, quantity,
+ base_price, disc_price, charge,
+ quantity, base_price, discount};
+ std::vector<std::string> project_names = {
+ "l_returnflag", "l_linestatus", "sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price", "avg_disc"};
+ ProjectNodeOptions project_opts(std::move(projection_list));
+
+ ScalarAggregateOptions sum_opts = ScalarAggregateOptions::Defaults();
+ CountOptions count_opts(CountOptions::CountMode::ALL);
+ std::vector<arrow::compute::internal::Aggregate> aggs = {
+ {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts},
+ {"hash_sum", &sum_opts}, {"hash_mean", &sum_opts}, {"hash_mean", &sum_opts},
+ {"hash_mean", &sum_opts}, {"hash_count", &count_opts}};
+
+ std::vector<FieldRef> cols = {2, 3, 4, 5, 6, 7, 8, 2};
Review comment:
why is it nicer?
##########
File path: cpp/src/arrow/compute/kernels/vector_sort.cc
##########
@@ -893,7 +893,8 @@ class TableSorter {
TableSorter(ExecContext* ctx, uint64_t* indices_begin, uint64_t* indices_end,
const Table& table, const SortOptions& options)
- : ctx_(ctx),
Review comment:
not in any very detailed way, just that it was inside this constructor, and there was some shared_ptr being deleted. it also didn't always reproduce, but it stopped reproducing once i made status be initialized before it got passed around to other stuff.
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
Review comment:
ok, i'll change it
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
Review comment:
yes, and it matters a lot. these strings will be placed adjacent in memory (maybe spaced out for alignment) in the .rodata section of the binary. this means that if i load these strings for whatever reason (like generating comments!) it'll pull a bunch of them in within the same cache line, and while generating comments the whole block of strings will fit within cache.
what happens with `std::vector<std::string>`? before `main` is called, all of the constructors of static objects will be called. this means that we will allocate the vector on the heap. then, we'll do a million tiny allocations for these strings, and these strings will be copied from `.rodata` onto the heap. further, we have no idea where the allocator will put these strings, they could just be in random places in memory (i.e. terrible for cache).
so basically we end up with something that's worse in every way: double the memory gets used, cache performance is worse, startup time is longer
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
+ thread_local_data_.resize(num_threads);
+ for (ThreadLocalData& tld : thread_local_data_) {
+ constexpr int kMaxNumDistinctStrings = 5;
+ tld.string_indices.resize(kMaxNumDistinctStrings * batch_size_);
+ tld.rng.seed(seq);
+ }
+ part_rows_to_generate_ = static_cast<int64_t>(scale_factor_ * 200000);
+ }
+ return Status::OK();
+ }
+
+ int64_t part_batches_generated() const { return part_batches_generated_.load(); }
+
+ int64_t partsupp_batches_generated() const {
+ return partsupp_batches_generated_.load();
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartTypes, kPartNameMap, part_cols_);
+ }
+
+ Result<std::shared_ptr<Schema>> SetPartSuppOutputColumns(
+ const std::vector<std::string>& cols) {
+ return SetOutputColumns(cols, kPartsuppTypes, kPartsuppNameMap, partsupp_cols_);
+ }
+
+ Result<util::optional<ExecBatch>> NextPartBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (!part_output_queue_.empty()) {
+ ExecBatch batch = std::move(part_output_queue_.front());
+ part_output_queue_.pop();
+ return std::move(batch);
+ } else if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ if (!partsupp_cols_.empty()) {
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ {
+ std::lock_guard<std::mutex> guard(partsupp_output_queue_mutex_);
+ for (ExecBatch& eb : partsupp_results) {
+ partsupp_output_queue_.emplace(std::move(eb));
+ }
+ }
+ }
+ return ExecBatch::Make(std::move(part_result));
+ }
+
+ Result<util::optional<ExecBatch>> NextPartSuppBatch() {
+ size_t thread_index = thread_indexer_();
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ if (!partsupp_output_queue_.empty()) {
+ ExecBatch result = std::move(partsupp_output_queue_.front());
+ partsupp_output_queue_.pop();
+ return std::move(result);
+ }
+ }
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ if (part_rows_generated_ == part_rows_to_generate_) {
+ return util::nullopt;
+ } else {
+ tld.partkey_start = part_rows_generated_;
+ tld.part_to_generate =
+ std::min(batch_size_, part_rows_to_generate_ - part_rows_generated_);
+ part_rows_generated_ += tld.part_to_generate;
+ int64_t num_ps_batches = PartsuppBatchesToGenerate(thread_index);
+ part_batches_generated_.fetch_add(1);
+ partsupp_batches_generated_.fetch_add(num_ps_batches);
+ ARROW_DCHECK(part_rows_generated_ <= part_rows_to_generate_);
+ }
+ }
+ tld.part.resize(PART::kNumCols);
+ std::fill(tld.part.begin(), tld.part.end(), Datum());
+ RETURN_NOT_OK(InitPartsupp(thread_index));
+
+ for (int col : part_cols_) RETURN_NOT_OK(kPartGenerators[col](thread_index));
+ for (int col : partsupp_cols_) RETURN_NOT_OK(kPartsuppGenerators[col](thread_index));
+ if (!part_cols_.empty()) {
+ std::vector<Datum> part_result(part_cols_.size());
+ for (size_t i = 0; i < part_cols_.size(); i++) {
+ int col_idx = part_cols_[i];
+ part_result[i] = tld.part[col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch part_batch,
+ ExecBatch::Make(std::move(part_result)));
+ {
+ std::lock_guard<std::mutex> lock(part_output_queue_mutex_);
+ part_output_queue_.emplace(std::move(part_batch));
+ }
+ }
+ std::vector<ExecBatch> partsupp_results;
+ for (size_t ibatch = 0; ibatch < tld.partsupp.size(); ibatch++) {
+ std::vector<Datum> partsupp_result(partsupp_cols_.size());
+ for (size_t icol = 0; icol < partsupp_cols_.size(); icol++) {
+ int col_idx = partsupp_cols_[icol];
+ partsupp_result[icol] = tld.partsupp[ibatch][col_idx];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch eb, ExecBatch::Make(std::move(partsupp_result)));
+ partsupp_results.emplace_back(std::move(eb));
+ }
+ // Return the first batch, enqueue the rest.
+ {
+ std::lock_guard<std::mutex> lock(partsupp_output_queue_mutex_);
+ for (size_t i = 1; i < partsupp_results.size(); i++)
+ partsupp_output_queue_.emplace(std::move(partsupp_results[i]));
+ }
+ return std::move(partsupp_results[0]);
+ }
+
+ private:
+#define FOR_EACH_PART_COLUMN(F) \
+ F(P_PARTKEY) \
+ F(P_NAME) \
+ F(P_MFGR) \
+ F(P_BRAND) \
+ F(P_TYPE) \
+ F(P_SIZE) \
+ F(P_CONTAINER) \
+ F(P_RETAILPRICE) \
+ F(P_COMMENT)
+
+#define FOR_EACH_PARTSUPP_COLUMN(F) \
+ F(PS_PARTKEY) \
+ F(PS_SUPPKEY) \
+ F(PS_AVAILQTY) \
+ F(PS_SUPPLYCOST) \
+ F(PS_COMMENT)
+
+#define MAKE_ENUM(col) col,
+ struct PART {
+ enum {
+ FOR_EACH_PART_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+ struct PARTSUPP {
+ enum {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_ENUM) kNumCols,
+ };
+ };
+
+#define MAKE_STRING_MAP(col) {#col, PART::col},
+ const std::unordered_map<std::string, int> kPartNameMap = {
+ FOR_EACH_PART_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_STRING_MAP(col) {#col, PARTSUPP::col},
+ const std::unordered_map<std::string, int> kPartsuppNameMap = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_STRING_MAP)};
+#undef MAKE_STRING_MAP
+#define MAKE_FN_ARRAY(col) \
+ [this](size_t thread_index) { return this->col(thread_index); },
+ std::vector<GenerateColumnFn> kPartGenerators = {FOR_EACH_PART_COLUMN(MAKE_FN_ARRAY)};
+ std::vector<GenerateColumnFn> kPartsuppGenerators = {
+ FOR_EACH_PARTSUPP_COLUMN(MAKE_FN_ARRAY)};
+#undef MAKE_FN_ARRAY
+#undef FOR_EACH_LINEITEM_COLUMN
+#undef FOR_EACH_ORDERS_COLUMN
+
+ const std::vector<std::shared_ptr<DataType>> kPartTypes = {
+ int32(), utf8(), fixed_size_binary(25), fixed_size_binary(10),
+ utf8(), int32(), fixed_size_binary(10), decimal(12, 2),
+ utf8(),
+ };
+
+ const std::vector<std::shared_ptr<DataType>> kPartsuppTypes = {
+ int32(), int32(), int32(), decimal(12, 2), utf8(),
+ };
+
+ Status AllocatePartBatch(size_t thread_index, int column) {
+ ThreadLocalData& tld = thread_local_data_[thread_index];
+ ARROW_DCHECK(tld.part[column].kind() == Datum::NONE);
+ int32_t byte_width = arrow::internal::GetByteWidth(*kPartTypes[column]);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> buff,
+ AllocateBuffer(tld.part_to_generate * byte_width));
+ ArrayData ad(kPartTypes[column], tld.part_to_generate, {nullptr, std::move(buff)});
+ tld.part[column] = std::move(ad);
+ return Status::OK();
+ }
+
+ Status P_PARTKEY(size_t thread_index) {
Review comment:
i added a big one at the top
##########
File path: cpp/src/arrow/compute/exec/tpch_node.cc
##########
@@ -0,0 +1,3431 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/unreachable.h"
+
+#include <algorithm>
+#include <bitset>
+#include <cstring>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+namespace arrow {
+using internal::checked_cast;
+
+namespace compute {
+const char* NameParts[] = {
+ "almond", "antique", "aquamarine", "azure", "beige", "bisque",
+ "black", "blanched", "blue", "blush", "brown", "burlywood",
+ "burnished", "chartreuse", "chiffon", "chocolate", "coral", "cornflower",
+ "cornsilk", "cream", "cyan", "dark", "deep", "dim",
+ "dodger", "drab", "firebrick", "floral", "forest", "frosted",
+ "gainsboro", "ghost", "goldenrod", "green", "grey", "honeydew",
+ "hot", "indian", "ivory", "khaki", "lace", "lavender",
+ "lawn", "lemon", "light", "lime", "linen", "magenta",
+ "maroon", "medium", "metallic", "midnight", "mint", "misty",
+ "moccasin", "navajo", "navy", "olive", "orange", "orchid",
+ "pale", "papaya", "peach", "peru", "pink", "plum",
+ "powder", "puff", "purple", "red", "rose", "rosy",
+ "royal", "saddle", "salmon", "sandy", "seashell", "sienna",
+ "sky", "slate", "smoke", "snow", "spring", "steel",
+ "tan", "thistle", "tomato", "turquoise", "violet", "wheat",
+ "white", "yellow",
+};
+static constexpr size_t kNumNameParts = sizeof(NameParts) / sizeof(NameParts[0]);
+
+const char* Types_1[] = {
+ "STANDARD ", "SMALL ", "MEDIUM ", "LARGE ", "ECONOMY ", "PROMO ",
+};
+static constexpr size_t kNumTypes_1 = sizeof(Types_1) / sizeof(Types_1[0]);
+
+const char* Types_2[] = {
+ "ANODIZED ", "BURNISHED ", "PLATED ", "POLISHED ", "BRUSHED ",
+};
+static constexpr size_t kNumTypes_2 = sizeof(Types_2) / sizeof(Types_2[0]);
+
+const char* Types_3[] = {
+ "TIN", "NICKEL", "BRASS", "STEEL", "COPPER",
+};
+static constexpr size_t kNumTypes_3 = sizeof(Types_3) / sizeof(Types_3[0]);
+
+const char* Containers_1[] = {
+ "SM ", "LG ", "MD ", "JUMBO ", "WRAP ",
+};
+static constexpr size_t kNumContainers_1 = sizeof(Containers_1) / sizeof(Containers_1[0]);
+
+const char* Containers_2[] = {
+ "CASE", "BOX", "BAG", "JAR", "PKG", "PACK", "CAN", "DRUM",
+};
+static constexpr size_t kNumContainers_2 = sizeof(Containers_2) / sizeof(Containers_2[0]);
+
+const char* Segments[] = {
+ "AUTOMOBILE", "BUILDING", "FURNITURE", "MACHINERY", "HOUSEHOLD",
+};
+static constexpr size_t kNumSegments = sizeof(Segments) / sizeof(Segments[0]);
+
+const char* Priorities[] = {
+ "1-URGENT", "2-HIGH", "3-MEDIUM", "4-NOT SPECIFIED", "5-LOW",
+};
+static constexpr size_t kNumPriorities = sizeof(Priorities) / sizeof(Priorities[0]);
+
+const char* Instructions[] = {
+ "DELIVER IN PERSON",
+ "COLLECT COD",
+ "NONE",
+ "TAKE BACK RETURN",
+};
+static constexpr size_t kNumInstructions = sizeof(Instructions) / sizeof(Instructions[0]);
+
+const char* Modes[] = {
+ "REG AIR", "AIR", "RAIL", "SHIP", "TRUCK", "MAIL", "FOB",
+};
+static constexpr size_t kNumModes = sizeof(Modes) / sizeof(Modes[0]);
+
+const char* Nouns[] = {
+ "foxes ", "ideas ", "theodolites ", "pinto beans ", "instructions ",
+ "dependencies ", "excuses ", "platelets ", "asymptotes ", "courts ",
+ "dolphins ", "multipliers ", "sautemes ", "warthogs ", "frets ",
+ "dinos ", "attainments ", "somas ", "Tiresias '", "patterns ",
+ "forges ", "braids ", "hockey players ", "frays ", "warhorses ",
+ "dugouts ", "notomis ", "epitaphs ", "pearls ", "tithes ",
+ "waters ", "orbits ", "gifts ", "sheaves ", "depths ",
+ "sentiments ", "decoys ", "realms ", "pains ", "grouches ",
+ "escapades ",
+};
+static constexpr size_t kNumNouns = sizeof(Nouns) / sizeof(Nouns[0]);
+
+const char* Verbs[] = {
+ "sleep ", "wake ", "are ", "cajole ", "haggle ", "nag ", "use ",
+ "boost ", "affix ", "detect ", "integrate ", "maintain ", "nod ", "was ",
+ "lose ", "sublate ", "solve ", "thrash ", "promise ", "engage ", "hinder ",
+ "print ", "x-ray ", "breach ", "eat ", "grow ", "impress ", "mold ",
+ "poach ", "serve ", "run ", "dazzle ", "snooze ", "doze ", "unwind ",
+ "kindle ", "play ", "hang ", "believe ", "doubt ",
+};
+static constexpr size_t kNumVerbs = sizeof(Verbs) / sizeof(Verbs[0]);
+
+const char* Adjectives[] = {
+ "furious ", "sly ", "careful ", "blithe ", "quick ", "fluffy ", "slow ",
+ "quiet ", "ruthless ", "thin ", "close ", "dogged ", "daring ", "brave ",
+ "stealthy ", "permanent ", "enticing ", "idle ", "busy ", "regular ", "final ",
+ "ironic ", "even ", "bold ", "silent ",
+};
+static constexpr size_t kNumAdjectives = sizeof(Adjectives) / sizeof(Adjectives[0]);
+
+const char* Adverbs[] = {
+ "sometimes ", "always ", "never ", "furiously ", "slyly ", "carefully ",
+ "blithely ", "quickly ", "fluffily ", "slowly ", "quietly ", "ruthlessly ",
+ "thinly ", "closely ", "doggedly ", "daringly ", "bravely ", "stealthily ",
+ "permanently ", "enticingly ", "idly ", "busily ", "regularly ", "finally ",
+ "ironically ", "evenly ", "boldly ", "silently ",
+};
+static constexpr size_t kNumAdverbs = sizeof(Adverbs) / sizeof(Adverbs[0]);
+
+const char* Prepositions[] = {
+ "about ", "above ", "according to ", "across ", "after ", "against ",
+ "along ", "alongside of ", "among ", "around ", "at ", "atop ",
+ "before ", "behind ", "beneath ", "beside ", "besides ", "between ",
+ "beyond ", "beyond ", "by ", "despite ", "during ", "except ",
+ "for ", "from ", "in place of ", "inside ", "instead of ", "into ",
+ "near ", "of ", "on ", "outside ", "over ", "past ",
+ "since ", "through ", "throughout ", "to ", "toward ", "under ",
+ "until ", "up ", "upon ", "without ", "with ", "within ",
+};
+static constexpr size_t kNumPrepositions = sizeof(Prepositions) / sizeof(Prepositions[0]);
+
+const char* Auxiliaries[] = {
+ "do ",
+ "may ",
+ "might ",
+ "shall ",
+ "will ",
+ "would ",
+ "can ",
+ "could ",
+ "should ",
+ "ought to ",
+ "must ",
+ "will have to ",
+ "shall have to ",
+ "could have to ",
+ "should have to ",
+ "must have to ",
+ "need to ",
+ "try to ",
+};
+static constexpr size_t kNumAuxiliaries = sizeof(Auxiliaries) / sizeof(Auxiliaries[0]);
+
+const char* Terminators[] = {
+ ".", ";", ":", "?", "!", "--",
+};
+static constexpr size_t kNumTerminators = sizeof(Terminators) / sizeof(Terminators[0]);
+
+// The spec says to generate a 300 MB string according to a grammar. This is a
+// concurrent implementation of the generator. Each thread generates the text in
+// (up to) 8KB chunks of text. The generator maintains a cursor into the
+// 300 MB buffer. After generating the chunk, the cursor is incremented
+// to reserve space, and the chunk is memcpy-d in.
+// This text is used to generate the COMMENT columns. To generate a comment, the spec
+// says to pick a random length and a random offset into the 300 MB buffer (it does
+// not specify it should be word/sentence aligned), and that slice of text becomes
+// the comment.
+class TpchPseudotext {
+ public:
+ Status EnsureInitialized(random::pcg32_fast& rng);
+ Result<Datum> GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length, random::pcg32_fast& rng);
+
+ private:
+ bool GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices);
+ bool GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdjective(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePreposition(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateTerminator(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+ bool GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ bool GenerateSentence(int64_t& offset, random::pcg32_fast& rng, char* arr);
+
+ std::atomic<bool> done_ = {false};
+ int64_t generated_offset_{0};
+ std::mutex text_guard_;
+ std::unique_ptr<Buffer> text_;
+ static constexpr int64_t kChunkSize = 8192;
+ static constexpr int64_t kTextBytes = 300 * 1024 * 1024; // 300 MB
+};
+
+static TpchPseudotext g_text;
+
+Status TpchPseudotext::EnsureInitialized(random::pcg32_fast& rng) {
+ if (done_.load()) return Status::OK();
+
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (!text_) {
+ ARROW_ASSIGN_OR_RAISE(text_, AllocateBuffer(kTextBytes));
+ }
+ }
+ char* out = reinterpret_cast<char*>(text_->mutable_data());
+ char temp_buff[kChunkSize];
+
+ while (!done_.load()) {
+ int64_t known_valid_offset = 0;
+ int64_t try_offset = 0;
+ while (GenerateSentence(try_offset, rng, temp_buff)) known_valid_offset = try_offset;
+
+ bool last_one;
+ int64_t offset;
+ int64_t memcpy_size;
+ {
+ std::lock_guard<std::mutex> lock(text_guard_);
+ if (done_.load()) return Status::OK();
+ int64_t bytes_remaining = kTextBytes - generated_offset_;
+ memcpy_size = std::min(known_valid_offset, bytes_remaining);
+ offset = generated_offset_;
+ generated_offset_ += memcpy_size;
+ last_one = generated_offset_ == kTextBytes;
+ }
+ std::memcpy(out + offset, temp_buff, memcpy_size);
+ if (last_one) done_.store(true);
+ }
+ return Status::OK();
+}
+
+Result<Datum> TpchPseudotext::GenerateComments(size_t num_comments, size_t min_length,
+ size_t max_length,
+ random::pcg32_fast& rng) {
+ RETURN_NOT_OK(EnsureInitialized(rng));
+ std::uniform_int_distribution<size_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buffer,
+ AllocateBuffer(sizeof(int32_t) * (num_comments + 1)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buffer->mutable_data());
+ offsets[0] = 0;
+ for (size_t i = 1; i <= num_comments; i++)
+ offsets[i] = offsets[i - 1] + static_cast<int32_t>(length_dist(rng));
+
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> comment_buffer,
+ AllocateBuffer(offsets[num_comments]));
+ char* comments = reinterpret_cast<char*>(comment_buffer->mutable_data());
+ for (size_t i = 0; i < num_comments; i++) {
+ size_t length = offsets[i + 1] - offsets[i];
+ std::uniform_int_distribution<size_t> offset_dist(0, kTextBytes - length);
+ size_t offset_in_text = offset_dist(rng);
+ std::memcpy(comments + offsets[i], text_->data() + offset_in_text, length);
+ }
+ ArrayData ad(utf8(), num_comments,
+ {nullptr, std::move(offset_buffer), std::move(comment_buffer)});
+ return std::move(ad);
+}
+
+bool TpchPseudotext::GenerateWord(int64_t& offset, random::pcg32_fast& rng, char* arr,
+ const char** words, size_t num_choices) {
+ std::uniform_int_distribution<size_t> dist(0, num_choices - 1);
+ const char* word = words[dist(rng)];
+ size_t length = std::strlen(word);
+ if (offset + length > kChunkSize) return false;
+ std::memcpy(arr + offset, word, length);
+ offset += length;
+ return true;
+}
+
+bool TpchPseudotext::GenerateNoun(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Nouns, kNumNouns);
+}
+
+bool TpchPseudotext::GenerateVerb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Verbs, kNumVerbs);
+}
+
+bool TpchPseudotext::GenerateAdjective(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Adjectives, kNumAdjectives);
+}
+
+bool TpchPseudotext::GenerateAdverb(int64_t& offset, random::pcg32_fast& rng, char* arr) {
+ return GenerateWord(offset, rng, arr, Adverbs, kNumAdverbs);
+}
+
+bool TpchPseudotext::GeneratePreposition(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Prepositions, kNumPrepositions);
+}
+
+bool TpchPseudotext::GenerateAuxiliary(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ return GenerateWord(offset, rng, arr, Auxiliaries, kNumAuxiliaries);
+}
+
+bool TpchPseudotext::GenerateTerminator(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ bool result = GenerateWord(offset, rng, arr, Terminators, kNumTerminators);
+ // Swap the space with the terminator
+ if (result) std::swap(*(arr + offset - 2), *(arr + offset - 1));
+ return result;
+}
+
+bool TpchPseudotext::GenerateNounPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ const char* comma_space = ", ";
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateWord(--offset, rng, arr, &comma_space, 1);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAdverb(offset, rng, arr);
+ success &= GenerateAdjective(offset, rng, arr);
+ success &= GenerateNoun(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GenerateVerbPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 3);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateAuxiliary(offset, rng, arr);
+ success &= GenerateVerb(offset, rng, arr);
+ success &= GenerateAdverb(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 3 inclusive");
+ break;
+ }
+ return success;
+}
+
+bool TpchPseudotext::GeneratePrepositionalPhrase(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ const char* the_space = "the ";
+ bool success = true;
+ success &= GeneratePreposition(offset, rng, arr);
+ success &= GenerateWord(offset, rng, arr, &the_space, 1);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ return success;
+}
+
+bool TpchPseudotext::GenerateSentence(int64_t& offset, random::pcg32_fast& rng,
+ char* arr) {
+ std::uniform_int_distribution<size_t> dist(0, 4);
+ bool success = true;
+ switch (dist(rng)) {
+ case 0:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 1:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 2:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 3:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ case 4:
+ success &= GenerateNounPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateVerbPhrase(offset, rng, arr);
+ success &= GeneratePrepositionalPhrase(offset, rng, arr);
+ success &= GenerateTerminator(offset, rng, arr);
+ break;
+ default:
+ Unreachable("Random number should be between 0 and 5 inclusive");
+ break;
+ }
+ return success;
+}
+
+class TpchTableGenerator {
+ public:
+ using OutputBatchCallback = std::function<void(ExecBatch)>;
+ using FinishedCallback = std::function<void(int64_t)>;
+ using GenerateFn = std::function<Status(size_t)>;
+ using ScheduleCallback = std::function<Status(GenerateFn)>;
+ using AbortCallback = std::function<void()>;
+
+ virtual Status Init(std::vector<std::string> columns, float scale_factor,
+ int64_t batch_size) = 0;
+
+ virtual Status StartProducing(size_t num_threads, OutputBatchCallback output_callback,
+ FinishedCallback finished_callback,
+ ScheduleCallback schedule_callback) = 0;
+
+ bool Abort() {
+ bool expected = false;
+ return done_.compare_exchange_strong(expected, true);
+ }
+
+ virtual std::shared_ptr<Schema> schema() const = 0;
+
+ virtual ~TpchTableGenerator() = default;
+
+ protected:
+ std::atomic<bool> done_ = {false};
+ std::atomic<int64_t> batches_outputted_ = {0};
+};
+
+int GetNumDigits(int64_t x) {
+ // This if statement chain is for MAXIMUM SPEED
+ // Source:
+ // https://stackoverflow.com/questions/1068849/how-do-i-determine-the-number-of-digits-of-an-integer-in-c
+ ARROW_DCHECK(x >= 0);
+ if (x < 10ll) return 1;
+ if (x < 100ll) return 2;
+ if (x < 1000ll) return 3;
+ if (x < 10000ll) return 4;
+ if (x < 100000ll) return 5;
+ if (x < 1000000ll) return 6;
+ if (x < 10000000ll) return 7;
+ if (x < 100000000ll) return 8;
+ if (x < 1000000000ll) return 9;
+ if (x < 10000000000ll) return 10;
+ if (x < 100000000000ll) return 11;
+ if (x < 1000000000000ll) return 12;
+ if (x < 10000000000000ll) return 13;
+ if (x < 100000000000000ll) return 14;
+ if (x < 1000000000000000ll) return 15;
+ if (x < 10000000000000000ll) return 16;
+ if (x < 100000000000000000ll) return 17;
+ if (x < 1000000000000000000ll) return 18;
+ return -1;
+}
+
+void AppendNumberPaddedToNineDigits(char* out, int64_t x) {
+ // We do all of this to avoid calling snprintf, which needs to handle locale,
+ // which can be slow, especially on Mac and Windows.
+ int num_digits = GetNumDigits(x);
+ int num_padding_zeros = std::max(9 - num_digits, 0);
+ std::memset(out, '0', static_cast<size_t>(num_padding_zeros));
+ while (x > 0) {
+ *(out + num_padding_zeros + num_digits - 1) = ('0' + x % 10);
+ num_digits -= 1;
+ x /= 10;
+ }
+}
+
+Result<std::shared_ptr<Schema>> SetOutputColumns(
+ const std::vector<std::string>& columns,
+ const std::vector<std::shared_ptr<DataType>>& types,
+ const std::unordered_map<std::string, int>& name_map, std::vector<int>& gen_list) {
+ gen_list.clear();
+ std::vector<std::shared_ptr<Field>> fields;
+ if (columns.empty()) {
+ fields.resize(name_map.size());
+ gen_list.resize(name_map.size());
+ for (auto pair : name_map) {
+ int col_idx = pair.second;
+ fields[col_idx] = field(pair.first, types[col_idx]);
+ gen_list[col_idx] = col_idx;
+ }
+ return schema(std::move(fields));
+ } else {
+ for (const std::string& col : columns) {
+ auto entry = name_map.find(col);
+ if (entry == name_map.end()) return Status::Invalid("Not a valid column name");
+ int col_idx = static_cast<int>(entry->second);
+ fields.push_back(field(col, types[col_idx]));
+ gen_list.push_back(col_idx);
+ }
+ return schema(std::move(fields));
+ }
+}
+
+Result<Datum> RandomVString(random::pcg32_fast& rng, int64_t num_rows, int32_t min_length,
+ int32_t max_length) {
+ std::uniform_int_distribution<int32_t> length_dist(min_length, max_length);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> offset_buff,
+ AllocateBuffer((num_rows + 1) * sizeof(int32_t)));
+ int32_t* offsets = reinterpret_cast<int32_t*>(offset_buff->mutable_data());
+ offsets[0] = 0;
+ for (int64_t i = 1; i <= num_rows; i++) offsets[i] = offsets[i - 1] + length_dist(rng);
+ ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Buffer> str_buff,
+ AllocateBuffer(offsets[num_rows]));
+ char* str = reinterpret_cast<char*>(str_buff->mutable_data());
+
+ // Spec says to pick random alphanumeric characters from a set of at least
+ // 64 symbols. Now, let's think critically here: 26 letters in the alphabet,
+ // so 52 total for upper and lower case, and 10 possible digits gives 62
+ // characters...
+ // dbgen solves this by including a space and a comma as well, so we'll
+ // copy that.
+ const char alpha_numerics[65] =
+ "0123456789abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ,";
+ std::uniform_int_distribution<int> char_dist(0, 63);
+ for (int32_t i = 0; i < offsets[num_rows]; i++) str[i] = alpha_numerics[char_dist(rng)];
+
+ ArrayData ad(utf8(), num_rows, {nullptr, std::move(offset_buff), std::move(str_buff)});
+ return std::move(ad);
+}
+
+void AppendNumber(char*& out, int num_digits, int32_t x) {
+ out += (num_digits - 1);
+ while (x > 0) {
+ *out-- = '0' + (x % 10);
+ x /= 10;
+ }
+ out += (num_digits + 1);
+}
+
+void GeneratePhoneNumber(char* out, random::pcg32_fast& rng, int32_t country) {
+ std::uniform_int_distribution<int32_t> three_digit(100, 999);
+ std::uniform_int_distribution<int32_t> four_digit(1000, 9999);
+
+ int32_t country_code = country + 10;
+ int32_t l1 = three_digit(rng);
+ int32_t l2 = three_digit(rng);
+ int32_t l3 = four_digit(rng);
+ AppendNumber(out, 2, country_code);
+ *out++ = '-';
+ AppendNumber(out, 3, l1);
+ *out++ = '-';
+ AppendNumber(out, 3, l2);
+ *out++ = '-';
+ AppendNumber(out, 4, l3);
+}
+
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kCurrentDate =
+ 9298; // June 17, 1995 is 9298 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+using GenerateColumnFn = std::function<Status(size_t)>;
+class PartAndPartSupplierGenerator {
+ public:
+ Status Init(size_t num_threads, int64_t batch_size, float scale_factor) {
+ if (!inited_) {
+ inited_ = true;
+ batch_size_ = batch_size;
+ scale_factor_ = scale_factor;
+
+ arrow_vendored::pcg_extras::seed_seq_from<std::random_device> seq;
Review comment:
i made it take a seed, defaulting to generating a nondeterministic one.
##########
File path: cpp/src/arrow/compute/exec/tpch_node_test.cc
##########
@@ -0,0 +1,623 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gmock/gmock-matchers.h>
+
+#include "arrow/api.h"
+#include "arrow/array/validate.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/compute/kernels/row_encoder.h"
+#include "arrow/compute/kernels/test_util.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/make_unique.h"
+#include "arrow/util/pcg_random.h"
+#include "arrow/util/thread_pool.h"
+
+#include <cctype>
+#include <string>
+#include <unordered_set>
+
+namespace arrow {
+namespace compute {
+static constexpr uint32_t kStartDate =
+ 8035; // January 1, 1992 is 8035 days after January 1, 1970
+static constexpr uint32_t kEndDate =
+ 10591; // December 12, 1998 is 10591 days after January 1, 1970
+
+void ValidateBatch(const ExecBatch& batch) {
+ for (const Datum& d : batch.values)
+ ASSERT_OK(arrow::internal::ValidateArray(*d.array()));
+}
+
+void VerifyUniqueKey(std::unordered_set<int32_t>& seen, const Datum& d, int32_t min,
+ int32_t max) {
+ const int32_t* keys = reinterpret_cast<const int32_t*>(d.array()->buffers[1]->data());
+ int64_t num_keys = d.length();
+ for (int64_t i = 0; i < num_keys; i++) {
+ ASSERT_TRUE(seen.insert(keys[i]).second);
+ ASSERT_LE(keys[i], max);
+ ASSERT_GE(keys[i], min);
+ }
+}
+
+void VerifyStringAndNumber_Single(const char* row, const char* prefix, const int64_t i,
+ const int32_t* nums, int byte_width,
+ bool verify_padding) {
+ int num_offset = static_cast<int>(std::strlen(prefix));
+ ASSERT_EQ(std::memcmp(row, prefix, num_offset), 0)
+ << row << ", prefix=" << prefix << ", i=" << i;
+ const char* num_str = row + num_offset;
+ int64_t num = 0;
+ int ibyte = static_cast<int>(num_offset);
+ for (; *num_str && ibyte < byte_width; ibyte++) {
+ num *= 10;
+ ASSERT_TRUE(std::isdigit(*num_str));
+ num += *num_str++ - '0';
+ }
+ if (nums) {
+ ASSERT_EQ(static_cast<int32_t>(num), nums[i]);
+ }
+ if (verify_padding) {
+ int num_chars = ibyte - num_offset;
+ ASSERT_GE(num_chars, 9);
+ }
+}
+
+void VerifyStringAndNumber_FixedWidth(const Datum& strings, const Datum& numbers,
+ int byte_width, const char* prefix,
+ bool verify_padding = true) {
+ int64_t length = strings.length();
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[1]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ const char* row = str + i * byte_width;
+ VerifyStringAndNumber_Single(row, prefix, i, nums, byte_width, verify_padding);
+ }
+}
+
+void VerifyStringAndNumber_Varlen(const Datum& strings, const Datum& numbers,
+ const char* prefix, bool verify_padding = true) {
+ int64_t length = strings.length();
+ const int32_t* offsets =
+ reinterpret_cast<const int32_t*>(strings.array()->buffers[1]->data());
+ const char* str = reinterpret_cast<const char*>(strings.array()->buffers[2]->data());
+
+ const int32_t* nums = nullptr;
+ if (numbers.kind() != Datum::NONE) {
+ ASSERT_EQ(length, numbers.length());
+ nums = reinterpret_cast<const int32_t*>(numbers.array()->buffers[1]->data());
+ }
+
+ for (int64_t i = 0; i < length; i++) {
+ char tmp_str[256] = {};
+ int32_t start = offsets[i];
+ int32_t str_len = offsets[i + 1] - offsets[i];
+ std::memcpy(tmp_str, str + start, str_len);
+ VerifyStringAndNumber_Single(tmp_str, prefix, i, nums, sizeof(tmp_str),
+ verify_padding);
+ }
Review comment:
i made it verify that there are no nulls when we verify the data.
regarding visitors, they look clunky and the code that implements it isn't very readable. i'd prefer to stick with the simple solution.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on a change in pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on a change in pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#discussion_r834592106
##########
File path: cpp/src/arrow/compute/exec/tpch_benchmark.cc
##########
@@ -0,0 +1,123 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <benchmark/benchmark.h>
+
+#include "arrow/compute/cast.h"
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/tpch_node.h"
+#include "arrow/testing/future_util.h"
+#include "arrow/util/make_unique.h"
+
+namespace arrow {
+namespace compute {
+namespace internal {
+
+std::shared_ptr<ExecPlan> Plan_Q1(AsyncGenerator<util::optional<ExecBatch>>* sink_gen,
+ int scale_factor) {
+ ExecContext* ctx = default_exec_context();
+ *ctx = ExecContext(default_memory_pool(), arrow::internal::GetCpuThreadPool());
+ std::shared_ptr<ExecPlan> plan = *ExecPlan::Make(ctx);
+ TpchGen gen = *TpchGen::Make(plan.get(), static_cast<double>(scale_factor));
+
+ ExecNode* lineitem =
+ *gen.Lineitem({"L_QUANTITY", "L_EXTENDEDPRICE", "L_TAX", "L_DISCOUNT", "L_SHIPDATE",
+ "L_RETURNFLAG", "L_LINESTATUS"});
+
+ auto sept_2_1998 = std::make_shared<Date32Scalar>(
+ 10471); // September 2, 1998 is 10471 days after January 1, 1970
+ Expression filter =
+ less_equal(field_ref("L_SHIPDATE"), literal(std::move(sept_2_1998)));
+ FilterNodeOptions filter_opts(filter);
+
+ Expression l_returnflag = field_ref("L_RETURNFLAG");
+ Expression l_linestatus = field_ref("L_LINESTATUS");
+ Expression quantity = field_ref("L_QUANTITY");
+ Expression base_price = field_ref("L_EXTENDEDPRICE");
+
+ std::shared_ptr<Decimal128Scalar> decimal_1 =
+ std::make_shared<Decimal128Scalar>(Decimal128{0, 100}, decimal(12, 2));
+ Expression discount_multiplier =
+ call("subtract", {literal(decimal_1), field_ref("L_DISCOUNT")});
+ Expression tax_multiplier = call("add", {literal(decimal_1), field_ref("L_TAX")});
+ Expression disc_price =
+ call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier});
+ Expression charge =
+ call("multiply",
+ {call("cast",
+ {call("multiply", {field_ref("L_EXTENDEDPRICE"), discount_multiplier})},
+ compute::CastOptions::Unsafe(decimal(12, 2))),
+ tax_multiplier});
+ Expression discount = field_ref("L_DISCOUNT");
+
+ std::vector<Expression> projection_list = {l_returnflag, l_linestatus, quantity,
+ base_price, disc_price, charge,
+ quantity, base_price, discount};
+ std::vector<std::string> project_names = {
+ "l_returnflag", "l_linestatus", "sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price", "avg_disc"};
+ ProjectNodeOptions project_opts(std::move(projection_list), std::move(project_names));
+
+ ScalarAggregateOptions sum_opts = ScalarAggregateOptions::Defaults();
+ CountOptions count_opts(CountOptions::CountMode::ALL);
+ std::vector<arrow::compute::internal::Aggregate> aggs = {
+ {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts}, {"hash_sum", &sum_opts},
+ {"hash_sum", &sum_opts}, {"hash_mean", &sum_opts}, {"hash_mean", &sum_opts},
+ {"hash_mean", &sum_opts}, {"hash_count", &count_opts}};
+
+ std::vector<FieldRef> to_aggregate = {"sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price",
+ "avg_disc", "sum_qty"};
+
+ std::vector<std::string> names = {"sum_qty", "sum_base_price", "sum_disc_price",
+ "sum_charge", "avg_qty", "avg_price",
+ "avg_disc", "count_order"};
+
+ std::vector<FieldRef> keys = {"l_returnflag", "l_linestatus"};
+ AggregateNodeOptions agg_opts(aggs, to_aggregate, names, keys);
+
+ SortKey l_returnflag_key("l_returnflag");
+ SortKey l_linestatus_key("l_linestatus");
+ SortOptions sort_opts({l_returnflag_key, l_linestatus_key});
+ OrderBySinkNodeOptions order_by_opts(sort_opts, sink_gen);
+
+ Declaration filter_decl("filter", {Declaration::Input(lineitem)}, filter_opts);
+ Declaration project_decl("project", project_opts);
+ Declaration aggregate_decl("aggregate", agg_opts);
+ Declaration orderby_decl("order_by_sink", order_by_opts);
+
+ Declaration q1 =
+ Declaration::Sequence({filter_decl, project_decl, aggregate_decl, orderby_decl});
+ std::ignore = *q1.AddToPlan(plan.get());
+ return plan;
+}
+
+static void BM_Tpch_Q1(benchmark::State& st) {
+ for (auto _ : st) {
+ st.PauseTiming();
+ AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+ std::shared_ptr<ExecPlan> plan = Plan_Q1(&sink_gen, static_cast<int>(st.range(0)));
+ st.ResumeTiming();
+ auto fut = StartAndCollect(plan.get(), sink_gen);
+ auto res = *fut.MoveResult();
+ }
+}
Review comment:
Q1 only returns four rows ever, so items per second might not be super useful, unless you mean input rows. however, Weston is creating a separate tool anyway so it this benchmark will be deleted very soon anyway in favor of that testing tool, so i don't think it's worth it adding more metrics to this one.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] pitrou removed a comment on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
pitrou removed a comment on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1077334837
I get the following warning when compiling the benchmarks:
```
[241/415] Building CXX object src/arrow/util/CMakeFiles/arrow-small-vector-benchmark.dir/small_vector_benchmark.cc.o
/home/antoine/arrow/dev/cpp/src/arrow/util/small_vector_benchmark.cc:70:37: warning: address of stack memory associated with parameter 'v' returned [-Wreturn-stack-address]
return reinterpret_cast<intptr_t>(v.data());
^
/home/antoine/arrow/dev/cpp/src/arrow/util/small_vector_benchmark.cc:119:16: note: in instantiation of function template specialization 'arrow::internal::ConsumeVector<std::vector<int>>' requested here
dummy += ConsumeVector(vec);
^
/home/antoine/arrow/dev/cpp/src/arrow/util/small_vector_benchmark.cc:325:1: note: in instantiation of function template specialization 'arrow::internal::CopyEmptyVector<std::vector<int>>' requested here
SHORT_VECTOR_BENCHMARKS(STD_VECTOR);
^
/home/antoine/arrow/dev/cpp/src/arrow/util/small_vector_benchmark.cc:287:22: note: expanded from macro 'SHORT_VECTOR_BENCHMARKS'
BENCHMARK_TEMPLATE(CopyEmptyVector, VEC_TYPE_FACTORY(int)); \
^
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] save-buffer commented on pull request #12537: ARROW-3998: [C++] Add TPC-H Generator
Posted by GitBox <gi...@apache.org>.
save-buffer commented on pull request #12537:
URL: https://github.com/apache/arrow/pull/12537#issuecomment-1078665714
Also I believe the spurious failure should be fixed. I had to add an extra method to `AsyncTaskGroup`, so please take a look at that.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org