parser.cpp 6.17 KB
Newer Older
1
2
3
4
/*!
 * Copyright (c) 2016 Microsoft Corporation. All rights reserved.
 * Licensed under the MIT License. See LICENSE file in the project root for license information.
 */
Guolin Ke's avatar
Guolin Ke committed
5
6
#include "parser.hpp"

7
#include <string>
Guolin Ke's avatar
Guolin Ke committed
8
#include <fstream>
Guolin Ke's avatar
Guolin Ke committed
9
#include <functional>
10
#include <iostream>
Guolin Ke's avatar
Guolin Ke committed
11
#include <memory>
Guolin Ke's avatar
Guolin Ke committed
12
13
14

namespace LightGBM {

Guolin Ke's avatar
Guolin Ke committed
15
void GetStatistic(const char* str, int* comma_cnt, int* tab_cnt, int* colon_cnt) {
Guolin Ke's avatar
Guolin Ke committed
16
17
18
19
20
21
22
23
24
25
26
27
28
29
  *comma_cnt = 0;
  *tab_cnt = 0;
  *colon_cnt = 0;
  for (int i = 0; str[i] != '\0'; ++i) {
    if (str[i] == ',') {
      ++(*comma_cnt);
    } else if (str[i] == '\t') {
      ++(*tab_cnt);
    } else if (str[i] == ':') {
      ++(*colon_cnt);
    }
  }
}

Guolin Ke's avatar
Guolin Ke committed
30
int GetLabelIdxForLibsvm(const std::string& str, int num_features, int label_idx) {
Guolin Ke's avatar
Guolin Ke committed
31
32
33
  if (num_features <= 0) {
    return label_idx;
  }
Guolin Ke's avatar
Guolin Ke committed
34
35
36
  auto str2 = Common::Trim(str);
  auto pos_space = str2.find_first_of(" \f\n\r\t\v");
  auto pos_colon = str2.find_first_of(":");
37
  if (pos_space == std::string::npos || pos_space < pos_colon) {
Guolin Ke's avatar
Guolin Ke committed
38
    return label_idx;
39
40
  } else {
    return -1;
41
42
43
  }
}

Guolin Ke's avatar
Guolin Ke committed
44
int GetLabelIdxForTSV(const std::string& str, int num_features, int label_idx) {
Guolin Ke's avatar
Guolin Ke committed
45
46
47
  if (num_features <= 0) {
    return label_idx;
  }
Guolin Ke's avatar
Guolin Ke committed
48
49
  auto str2 = Common::Trim(str);
  auto tokens = Common::Split(str2.c_str(), '\t');
Guolin Ke's avatar
Guolin Ke committed
50
  if (static_cast<int>(tokens.size()) == num_features) {
Guolin Ke's avatar
Guolin Ke committed
51
    return -1;
52
  } else {
Guolin Ke's avatar
Guolin Ke committed
53
    return label_idx;
54
55
56
  }
}

Guolin Ke's avatar
Guolin Ke committed
57
int GetLabelIdxForCSV(const std::string& str, int num_features, int label_idx) {
Guolin Ke's avatar
Guolin Ke committed
58
59
60
  if (num_features <= 0) {
    return label_idx;
  }
Guolin Ke's avatar
Guolin Ke committed
61
62
  auto str2 = Common::Trim(str);
  auto tokens = Common::Split(str2.c_str(), ',');
Guolin Ke's avatar
Guolin Ke committed
63
  if (static_cast<int>(tokens.size()) == num_features) {
Guolin Ke's avatar
Guolin Ke committed
64
    return -1;
65
  } else {
Guolin Ke's avatar
Guolin Ke committed
66
    return label_idx;
67
68
69
  }
}

Guolin Ke's avatar
Guolin Ke committed
70
71
72
73
74
75
76
enum DataType {
  INVALID,
  CSV,
  TSV,
  LIBSVM
};

Guolin Ke's avatar
Guolin Ke committed
77
78
79
80
void GetLine(std::stringstream* ss, std::string* line, const VirtualFileReader* reader, std::vector<char>* buffer, size_t buffer_size) {
  std::getline(*ss, *line);
  while (ss->eof()) {
    size_t read_len = reader->Read(buffer->data(), buffer_size);
81
82
83
    if (read_len <= 0) {
      break;
    }
Guolin Ke's avatar
Guolin Ke committed
84
85
    ss->clear();
    ss->str(std::string(buffer->data(), read_len));
86
    std::string tmp;
Guolin Ke's avatar
Guolin Ke committed
87
88
    std::getline(*ss, tmp);
    *line += tmp;
89
90
91
  }
}

92
std::vector<std::string> ReadKLineFromFile(const char* filename, bool header, int k) {
93
94
  auto reader = VirtualFileReader::Make(filename);
  if (!reader->Init()) {
95
    Log::Fatal("Data file %s doesn't exist.", filename);
Guolin Ke's avatar
Guolin Ke committed
96
  }
97
98
99
  std::vector<std::string> ret;
  std::string cur_line;
  const size_t buffer_size = 1024 * 1024;
100
101
102
  auto buffer = std::vector<char>(buffer_size);
  size_t read_len = reader->Read(buffer.data(), buffer_size);
  if (read_len <= 0) {
103
    Log::Fatal("Data file %s couldn't be read.", filename);
104
  }
105
106
  std::string read_str = std::string(buffer.data(), read_len);
  std::stringstream tmp_file(read_str);
Guolin Ke's avatar
Guolin Ke committed
107
  if (header) {
Guolin Ke's avatar
Guolin Ke committed
108
    if (!tmp_file.eof()) {
109
      GetLine(&tmp_file, &cur_line, reader.get(), &buffer, buffer_size);
Guolin Ke's avatar
Guolin Ke committed
110
111
    }
  }
112
113
114
115
116
117
118
  for (int i = 0; i < k; ++i) {
    if (!tmp_file.eof()) {
      GetLine(&tmp_file, &cur_line, reader.get(), &buffer, buffer_size);
      ret.push_back(cur_line);
    } else {
      break;
    }
Guolin Ke's avatar
Guolin Ke committed
119
  }
120
121
122
123
124
125
126
  if (ret.empty()) {
    Log::Fatal("Data file %s should have at least one line.", filename);
  } else if (ret.size() == 1) {
    Log::Warning("Data file %s only has one line.", filename);
  }
  return ret;
}
Guolin Ke's avatar
Guolin Ke committed
127

128
DataType GetDataType(const std::vector<std::string>& lines, int* num_col) {
Guolin Ke's avatar
Guolin Ke committed
129
  DataType type = DataType::INVALID;
130
131
132
133
134
135
136
137
  if (lines.empty()) {
    return type;
  }
  int comma_cnt = 0;
  int tab_cnt = 0;
  int colon_cnt = 0;
  GetStatistic(lines[0].c_str(), &comma_cnt, &tab_cnt, &colon_cnt);
  if (lines.size() == 1) {
Guolin Ke's avatar
Guolin Ke committed
138
    if (colon_cnt > 0) {
Guolin Ke's avatar
Guolin Ke committed
139
      type = DataType::LIBSVM;
Guolin Ke's avatar
Guolin Ke committed
140
    } else if (tab_cnt > 0) {
Guolin Ke's avatar
Guolin Ke committed
141
      type = DataType::TSV;
Guolin Ke's avatar
Guolin Ke committed
142
    } else if (comma_cnt > 0) {
Guolin Ke's avatar
Guolin Ke committed
143
144
      type = DataType::CSV;
    }
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
  }
  int comma_cnt2 = 0;
  int tab_cnt2 = 0;
  int colon_cnt2 = 0;
  GetStatistic(lines[1].c_str(), &comma_cnt2, &tab_cnt2, &colon_cnt2);
  if (colon_cnt > 0 || colon_cnt2 > 0) {
    type = DataType::LIBSVM;
  } else if (tab_cnt == tab_cnt2 && tab_cnt > 0) {
    type = DataType::TSV;
  } else if (comma_cnt == comma_cnt2 && comma_cnt > 0) {
    type = DataType::CSV;
  }
  if (type == DataType::TSV || type == DataType::CSV) {
    // valid the type
    for (size_t i = 2; i < lines.size(); ++i) {
      GetStatistic(lines[i].c_str(), &comma_cnt2, &tab_cnt2, &colon_cnt2);
      if (type == DataType::TSV && tab_cnt2 != tab_cnt) {
        type = DataType::INVALID;
        break;
      } else if (type == DataType::CSV && comma_cnt != comma_cnt2) {
        type = DataType::INVALID;
        break;
      }
    }
  }

  if (type == DataType::LIBSVM) {
    int max_col_idx = 0;
    for (size_t i = 0; i < lines.size(); ++i) {
      auto str = Common::Trim(lines[i]);
      auto colon_pos = str.find_last_of(":");
      auto space_pos = str.find_last_of(" \f\t\v");
      auto sub_str = str.substr(space_pos + 1, space_pos - colon_pos - 1);
      int cur_idx = 0;
      Common::Atoi(sub_str.c_str(), &cur_idx);
      max_col_idx = std::max(cur_idx, max_col_idx);
Guolin Ke's avatar
Guolin Ke committed
181
    }
182
183
184
185
186
    *num_col = max_col_idx + 1;
  } else if (type == DataType::CSV) {
    *num_col = comma_cnt + 1;
  } else if (type == DataType::TSV) {
    *num_col = tab_cnt + 1;
Guolin Ke's avatar
Guolin Ke committed
187
  }
188
189
190
191
192
193
194
195
  return type;
}

Parser* Parser::CreateParser(const char* filename, bool header, int num_features, int label_idx) {
  const int n_read_line = 20;
  auto lines = ReadKLineFromFile(filename, header, n_read_line);
  int num_col = 0;
  DataType type = GetDataType(lines, &num_col);
Guolin Ke's avatar
Guolin Ke committed
196
  if (type == DataType::INVALID) {
197
    Log::Fatal("Unknown format of training data.");
Guolin Ke's avatar
Guolin Ke committed
198
  }
Guolin Ke's avatar
Guolin Ke committed
199
  std::unique_ptr<Parser> ret;
Guolin Ke's avatar
Guolin Ke committed
200
  if (type == DataType::LIBSVM) {
201
202
    label_idx = GetLabelIdxForLibsvm(lines[0], num_features, label_idx);
    ret.reset(new LibSVMParser(label_idx, num_col));
203
  } else if (type == DataType::TSV) {
204
205
    label_idx = GetLabelIdxForTSV(lines[0], num_features, label_idx);
    ret.reset(new TSVParser(label_idx, num_col));
206
  } else if (type == DataType::CSV) {
207
208
    label_idx = GetLabelIdxForCSV(lines[0], num_features, label_idx);
    ret.reset(new CSVParser(label_idx, num_col));
Guolin Ke's avatar
Guolin Ke committed
209
210
211
  }

  if (label_idx < 0) {
212
    Log::Info("Data file %s doesn't contain a label column.", filename);
Guolin Ke's avatar
Guolin Ke committed
213
  }
Guolin Ke's avatar
Guolin Ke committed
214
  return ret.release();
Guolin Ke's avatar
Guolin Ke committed
215
216
217
}

}  // namespace LightGBM