| 819 | } |
| 820 | |
| 821 | std::vector<std::string> DatasetLoader::SampleTextDataFromFile(const char* filename, const Metadata& metadata, |
| 822 | int rank, int num_machines, int* num_global_data, |
| 823 | std::vector<data_size_t>* used_data_indices) { |
| 824 | const data_size_t sample_cnt = static_cast<data_size_t>(config_.bin_construct_sample_cnt); |
| 825 | TextReader<data_size_t> text_reader(filename, config_.header, config_.file_load_progress_interval_bytes); |
| 826 | std::vector<std::string> out_data; |
| 827 | if (num_machines == 1 || config_.pre_partition) { |
| 828 | *num_global_data = static_cast<data_size_t>(text_reader.SampleFromFile(&random_, sample_cnt, &out_data)); |
| 829 | } else { // need partition data |
| 830 | // get query data |
| 831 | const data_size_t* query_boundaries = metadata.query_boundaries(); |
| 832 | if (query_boundaries == nullptr) { |
| 833 | // if not contain query file, minimal sample unit is one record |
| 834 | *num_global_data = text_reader.SampleAndFilterFromFile([this, rank, num_machines] |
| 835 | (data_size_t) { |
| 836 | if (random_.NextShort(0, num_machines) == rank) { |
| 837 | return true; |
| 838 | } else { |
| 839 | return false; |
| 840 | } |
| 841 | }, used_data_indices, &random_, sample_cnt, &out_data); |
| 842 | } else { |
| 843 | // if contain query file, minimal sample unit is one query |
| 844 | data_size_t num_queries = metadata.num_queries(); |
| 845 | data_size_t qid = -1; |
| 846 | bool is_query_used = false; |
| 847 | *num_global_data = text_reader.SampleAndFilterFromFile( |
| 848 | [this, rank, num_machines, &qid, &query_boundaries, &is_query_used, num_queries] |
| 849 | (data_size_t line_idx) { |
| 850 | if (qid >= num_queries) { |
| 851 | Log::Fatal("Query id exceeds the range of the query file, " |
| 852 | "please ensure the query file is correct"); |
| 853 | } |
| 854 | if (line_idx >= query_boundaries[qid + 1]) { |
| 855 | // if is new query |
| 856 | is_query_used = false; |
| 857 | if (random_.NextShort(0, num_machines) == rank) { |
| 858 | is_query_used = true; |
| 859 | } |
| 860 | ++qid; |
| 861 | } |
| 862 | return is_query_used; |
| 863 | }, used_data_indices, &random_, sample_cnt, &out_data); |
| 864 | } |
| 865 | } |
| 866 | return out_data; |
| 867 | } |
| 868 | |
| 869 | void DatasetLoader::ConstructBinMappersFromTextData(int rank, int num_machines, |
| 870 | const std::vector<std::string>& sample_data, |
nothing calls this directly
no test coverage detected