]>
Commit | Line | Data |
---|---|---|
7c673cae | 1 | // Copyright (c) 2011-present, Facebook, Inc. All rights reserved. |
11fdf7f2 TL |
2 | // This source code is licensed under both the GPLv2 (found in the |
3 | // COPYING file in the root directory) and Apache 2.0 License | |
4 | // (found in the LICENSE.Apache file in the root directory). | |
7c673cae FG |
5 | // |
6 | #ifndef __STDC_FORMAT_MACROS | |
7 | #define __STDC_FORMAT_MACROS | |
8 | #endif | |
9 | ||
10 | #include <cstdio> | |
11 | #include <cstdlib> | |
12 | ||
13 | #ifndef ROCKSDB_LITE | |
14 | #ifdef GFLAGS | |
15 | ||
7c673cae FG |
16 | #include <inttypes.h> |
17 | #include <vector> | |
18 | #include "rocksdb/env.h" | |
19 | #include "rocksdb/options.h" | |
20 | #include "table/block_based_table_builder.h" | |
21 | #include "table/block_based_table_reader.h" | |
22 | #include "table/format.h" | |
23 | #include "tools/sst_dump_tool_imp.h" | |
24 | #include "util/compression.h" | |
11fdf7f2 | 25 | #include "util/gflags_compat.h" |
7c673cae FG |
26 | #include "util/stop_watch.h" |
27 | #include "utilities/col_buf_encoder.h" | |
28 | #include "utilities/column_aware_encoding_util.h" | |
29 | ||
11fdf7f2 | 30 | using GFLAGS_NAMESPACE::ParseCommandLineFlags; |
7c673cae FG |
31 | DEFINE_string(encoded_file, "", "file to store encoded data blocks"); |
32 | DEFINE_string(decoded_file, "", | |
33 | "file to store decoded data blocks after encoding"); | |
34 | DEFINE_string(format, "col", "Output Format. Can be 'row' or 'col'"); | |
35 | // TODO(jhli): option `col` should be removed and replaced by general | |
36 | // column specifications. | |
37 | DEFINE_string(index_type, "col", "Index type. Can be 'primary' or 'secondary'"); | |
38 | DEFINE_string(dump_file, "", | |
39 | "Dump data blocks separated by columns in human-readable format"); | |
40 | DEFINE_bool(decode, false, "Deocde blocks after they are encoded"); | |
41 | DEFINE_bool(stat, false, | |
42 | "Print column distribution statistics. Cannot decode in this mode"); | |
43 | DEFINE_string(compression_type, "kNoCompression", | |
44 | "The compression algorithm used to compress data blocks"); | |
45 | ||
46 | namespace rocksdb { | |
47 | ||
48 | class ColumnAwareEncodingExp { | |
49 | public: | |
50 | static void Run(const std::string& sst_file) { | |
51 | bool decode = FLAGS_decode; | |
52 | if (FLAGS_decoded_file.size() > 0) { | |
53 | decode = true; | |
54 | } | |
55 | if (FLAGS_stat) { | |
56 | decode = false; | |
57 | } | |
58 | ||
59 | ColumnAwareEncodingReader reader(sst_file); | |
60 | std::vector<ColDeclaration>* key_col_declarations; | |
61 | std::vector<ColDeclaration>* value_col_declarations; | |
62 | ColDeclaration* value_checksum_declaration; | |
63 | if (FLAGS_index_type == "primary") { | |
64 | ColumnAwareEncodingReader::GetColDeclarationsPrimary( | |
65 | &key_col_declarations, &value_col_declarations, | |
66 | &value_checksum_declaration); | |
67 | } else { | |
68 | ColumnAwareEncodingReader::GetColDeclarationsSecondary( | |
69 | &key_col_declarations, &value_col_declarations, | |
70 | &value_checksum_declaration); | |
71 | } | |
72 | KVPairColDeclarations kvp_cd(key_col_declarations, value_col_declarations, | |
73 | value_checksum_declaration); | |
74 | ||
75 | if (!FLAGS_dump_file.empty()) { | |
76 | std::vector<KVPairBlock> kv_pair_blocks; | |
77 | reader.GetKVPairsFromDataBlocks(&kv_pair_blocks); | |
78 | reader.DumpDataColumns(FLAGS_dump_file, kvp_cd, kv_pair_blocks); | |
79 | return; | |
80 | } | |
81 | std::unordered_map<std::string, CompressionType> compressions = { | |
82 | {"kNoCompression", CompressionType::kNoCompression}, | |
83 | {"kZlibCompression", CompressionType::kZlibCompression}, | |
84 | {"kZSTD", CompressionType::kZSTD}}; | |
85 | ||
86 | // Find Compression | |
87 | CompressionType compression_type = compressions[FLAGS_compression_type]; | |
88 | EnvOptions env_options; | |
89 | if (CompressionTypeSupported(compression_type)) { | |
90 | fprintf(stdout, "[%s]\n", FLAGS_compression_type.c_str()); | |
91 | unique_ptr<WritableFile> encoded_out_file; | |
92 | ||
93 | std::unique_ptr<Env> env(NewMemEnv(Env::Default())); | |
94 | if (!FLAGS_encoded_file.empty()) { | |
95 | env->NewWritableFile(FLAGS_encoded_file, &encoded_out_file, | |
96 | env_options); | |
97 | } | |
98 | ||
99 | std::vector<KVPairBlock> kv_pair_blocks; | |
100 | reader.GetKVPairsFromDataBlocks(&kv_pair_blocks); | |
101 | ||
102 | std::vector<std::string> encoded_blocks; | |
103 | StopWatchNano sw(env.get(), true); | |
104 | if (FLAGS_format == "col") { | |
105 | reader.EncodeBlocks(kvp_cd, encoded_out_file.get(), compression_type, | |
106 | kv_pair_blocks, &encoded_blocks, FLAGS_stat); | |
107 | } else { // row format | |
108 | reader.EncodeBlocksToRowFormat(encoded_out_file.get(), compression_type, | |
109 | kv_pair_blocks, &encoded_blocks); | |
110 | } | |
111 | if (encoded_out_file != nullptr) { | |
112 | uint64_t size = 0; | |
113 | env->GetFileSize(FLAGS_encoded_file, &size); | |
114 | fprintf(stdout, "File size: %" PRIu64 "\n", size); | |
115 | } | |
116 | uint64_t encode_time = sw.ElapsedNanosSafe(false /* reset */); | |
117 | fprintf(stdout, "Encode time: %" PRIu64 "\n", encode_time); | |
118 | if (decode) { | |
119 | unique_ptr<WritableFile> decoded_out_file; | |
120 | if (!FLAGS_decoded_file.empty()) { | |
121 | env->NewWritableFile(FLAGS_decoded_file, &decoded_out_file, | |
122 | env_options); | |
123 | } | |
124 | sw.Start(); | |
125 | if (FLAGS_format == "col") { | |
126 | reader.DecodeBlocks(kvp_cd, decoded_out_file.get(), &encoded_blocks); | |
127 | } else { | |
128 | reader.DecodeBlocksFromRowFormat(decoded_out_file.get(), | |
129 | &encoded_blocks); | |
130 | } | |
131 | uint64_t decode_time = sw.ElapsedNanosSafe(true /* reset */); | |
132 | fprintf(stdout, "Decode time: %" PRIu64 "\n", decode_time); | |
133 | } | |
134 | } else { | |
135 | fprintf(stdout, "Unsupported compression type: %s.\n", | |
136 | FLAGS_compression_type.c_str()); | |
137 | } | |
138 | delete key_col_declarations; | |
139 | delete value_col_declarations; | |
140 | delete value_checksum_declaration; | |
141 | } | |
142 | }; | |
143 | ||
144 | } // namespace rocksdb | |
145 | ||
146 | int main(int argc, char** argv) { | |
147 | int arg_idx = ParseCommandLineFlags(&argc, &argv, true); | |
148 | if (arg_idx >= argc) { | |
149 | fprintf(stdout, "SST filename required.\n"); | |
150 | exit(1); | |
151 | } | |
152 | std::string sst_file(argv[arg_idx]); | |
153 | if (FLAGS_format != "row" && FLAGS_format != "col") { | |
154 | fprintf(stderr, "Format must be 'row' or 'col'\n"); | |
155 | exit(1); | |
156 | } | |
157 | if (FLAGS_index_type != "primary" && FLAGS_index_type != "secondary") { | |
158 | fprintf(stderr, "Format must be 'primary' or 'secondary'\n"); | |
159 | exit(1); | |
160 | } | |
161 | rocksdb::ColumnAwareEncodingExp::Run(sst_file); | |
162 | return 0; | |
163 | } | |
164 | ||
165 | #else | |
166 | int main() { | |
167 | fprintf(stderr, "Please install gflags to run rocksdb tools\n"); | |
168 | return 1; | |
169 | } | |
170 | #endif // GFLAGS | |
171 | #else | |
11fdf7f2 | 172 | int main(int /*argc*/, char** /*argv*/) { |
7c673cae FG |
173 | fprintf(stderr, "Not supported in lite mode.\n"); |
174 | return 1; | |
175 | } | |
176 | #endif // ROCKSDB_LITE |