diff options
author | Vadim Averin <avevad0808@gmail.com> | 2024-06-18 15:12:30 +0300 |
---|---|---|
committer | GitHub <noreply@github.com> | 2024-06-18 15:12:30 +0300 |
commit | 16145e3aa5e38e603c3ab43cc432c3ab3a9b0c9e (patch) | |
tree | c3fb1344262bbe45ef11a0e12806c230b8d5bd0f | |
parent | 6ff66f609fef8bf83c298d1fb0c47868631ad3ae (diff) | |
download | ydb-16145e3aa5e38e603c3ab43cc432c3ab3a9b0c9e.tar.gz |
Add support for RANGE and FOLDER in yqlrun (#5407)
Co-authored-by: Vadim Averin <avevad@ydb.tech>
16 files changed, 117 insertions, 66 deletions
diff --git a/ydb/library/yql/providers/yt/gateway/file/yql_yt_file.cpp b/ydb/library/yql/providers/yt/gateway/file/yql_yt_file.cpp index 4a44d0449f..0df77e9284 100644 --- a/ydb/library/yql/providers/yt/gateway/file/yql_yt_file.cpp +++ b/ydb/library/yql/providers/yt/gateway/file/yql_yt_file.cpp @@ -453,18 +453,26 @@ public: auto pos = options.Pos(); try { TSession* session = GetSession(options); + TSet<TString> uniqueTables; - if (options.Prefix().empty() && options.Suffix().empty()) { - for (auto& x : Services_->GetTablesMapping()) { - TVector<TString> parts; - Split(x.first, ".", parts); - if (parts.size() > 2 && parts[0] == YtProviderName) { - if (!parts[2].StartsWith(TStringBuf("Input"))) { - continue; - } - uniqueTables.insert(parts[2]); - } + const auto fullPrefix = options.Prefix().Empty() ? TString() : (options.Prefix() + '/'); + const auto fullSuffix = options.Suffix().Empty() ? TString() : ('/' + options.Suffix()); + for (const auto& [tableName, _] : Services_->GetTablesMapping()) { + TVector<TString> parts; + Split(tableName, ".", parts); + if (parts.size() != 3) { + continue; + } + if (parts[0] != YtProviderName || parts[1] != options.Cluster()) { + continue; + } + if (!parts[2].StartsWith(fullPrefix)) { + continue; + } + if (!parts[2].EndsWith(fullSuffix)) { + continue; } + uniqueTables.insert(parts[2]); } TTableRangeResult res; @@ -484,8 +492,11 @@ public: TProgramBuilder pgmBuilder(builder.GetTypeEnvironment(), *Services_->GetFunctionRegistry()); TVector<TRuntimeNode> strings; - for (auto& x: uniqueTables) { - strings.push_back(pgmBuilder.NewDataLiteral<NUdf::EDataSlot::String>(x)); + for (auto& tableName: uniqueTables) { + auto stripped = TStringBuf(tableName); + stripped.SkipPrefix(fullPrefix); + stripped.ChopSuffix(fullSuffix); + strings.push_back(pgmBuilder.NewDataLiteral<NUdf::EDataSlot::String>(TString(stripped))); } auto inputNode = pgmBuilder.AsList(strings); @@ -504,7 +515,10 @@ public: const auto& value = compGraph->GetValue(); const auto it = value.GetListIterator(); for (NUdf::TUnboxedValue current; it.Next(current);) { - res.Tables.push_back(TCanonizedPath{TString(current.AsStringRef()), Nothing(), {}, Nothing()}); + TString tableName = TString(current.AsStringRef()); + tableName.prepend(fullPrefix); + tableName.append(fullSuffix); + res.Tables.push_back(TCanonizedPath{std::move(tableName), Nothing(), {}, Nothing()}); } } else { @@ -527,17 +541,20 @@ public: auto pos = options.Pos(); try { TSet<TString> uniqueTables; - if (options.Prefix().empty()) { - for (auto& x : Services_->GetTablesMapping()) { - TVector<TString> parts; - Split(x.first, ".", parts); - if (parts.size() > 2 && parts[0] == YtProviderName) { - if (!parts[2].StartsWith(TStringBuf("Input"))) { - continue; - } - uniqueTables.insert(parts[2]); - } + const auto fullPrefix = options.Prefix().Empty() ? "" : (options.Prefix() + '/'); + for (const auto& [tableName, _] : Services_->GetTablesMapping()) { + TVector<TString> parts; + Split(tableName, ".", parts); + if (parts.size() != 3) { + continue; + } + if (parts[0] != YtProviderName || parts[1] != options.Cluster()) { + continue; + } + if (!parts[2].StartsWith(fullPrefix)) { + continue; } + uniqueTables.insert(parts[2]); } TVector<TFolderResult::TFolderItem> items; diff --git a/ydb/library/yql/providers/yt/gateway/file/yql_yt_file_services.cpp b/ydb/library/yql/providers/yt/gateway/file/yql_yt_file_services.cpp index c26d3425d8..7850bace15 100644 --- a/ydb/library/yql/providers/yt/gateway/file/yql_yt_file_services.cpp +++ b/ydb/library/yql/providers/yt/gateway/file/yql_yt_file_services.cpp @@ -31,7 +31,8 @@ TString TYtFileServices::GetTablePath(TStringBuf cluster, TStringBuf table, bool return TString(TFsPath(TmpDir) / TString(table.substr(4)).append(TStringBuf(".tmp"))); } - auto fullTableName = TString(YtProviderName).append('.').append(cluster).append('.').append(table); + const auto tablePrefix = TString(YtProviderName).append('.').append(cluster); + const auto fullTableName = TString(tablePrefix).append('.').append(table); if (!noLocks) { auto guard = Guard(Mutex); if (auto p = Locks.FindPtr(fullTableName)) { @@ -41,6 +42,9 @@ TString TYtFileServices::GetTablePath(TStringBuf cluster, TStringBuf table, bool if (auto p = TablesMapping.FindPtr(fullTableName)) { return *p; } + if (auto dirPtr = TablesDirMapping.FindPtr(tablePrefix)) { + return TFsPath(*dirPtr) / TString(table).append(".txt"); + } ythrow yexception() << "Table not found: " << cluster << '.' << table; } diff --git a/ydb/library/yql/providers/yt/gateway/file/yql_yt_file_services.h b/ydb/library/yql/providers/yt/gateway/file/yql_yt_file_services.h index 43cc413f19..c29fc64554 100644 --- a/ydb/library/yql/providers/yt/gateway/file/yql_yt_file_services.h +++ b/ydb/library/yql/providers/yt/gateway/file/yql_yt_file_services.h @@ -22,9 +22,9 @@ public: ~TYtFileServices(); static TPtr Make(const NKikimr::NMiniKQL::IFunctionRegistry* registry, const THashMap<TString, TString>& mapping = {}, - TFileStoragePtr fileStorage = {}, const TString& tmpDir = {}, bool keepTempTables = false) + TFileStoragePtr fileStorage = {}, const TString& tmpDir = {}, bool keepTempTables = false, const THashMap<TString, TString>& dirMapping = {}) { - return new TYtFileServices(registry, mapping, fileStorage, tmpDir.empty() ? GetSystemTempDir() : tmpDir, keepTempTables); + return new TYtFileServices(registry, mapping, fileStorage, tmpDir.empty() ? GetSystemTempDir() : tmpDir, keepTempTables, dirMapping); } const NKikimr::NMiniKQL::IFunctionRegistry* GetFunctionRegistry() const { @@ -55,9 +55,17 @@ public: } private: - TYtFileServices(const NKikimr::NMiniKQL::IFunctionRegistry* registry, const THashMap<TString, TString>& mapping, TFileStoragePtr fileStorage, const TString& tmpDir, bool keepTempTables) + TYtFileServices( + const NKikimr::NMiniKQL::IFunctionRegistry* registry, + const THashMap<TString, TString>& mapping, + TFileStoragePtr fileStorage, + const TString& tmpDir, + bool keepTempTables, + const THashMap<TString, TString>& dirMapping + ) : FunctionRegistry(registry) , TablesMapping(mapping) + , TablesDirMapping(dirMapping) , TmpDir(tmpDir) , KeepTempTables(keepTempTables) { @@ -71,6 +79,7 @@ private: TFileStoragePtr FileStorage; const NKikimr::NMiniKQL::IFunctionRegistry* FunctionRegistry; THashMap<TString, TString> TablesMapping; // [cluster].[name] -> [file path] + THashMap<TString, TString> TablesDirMapping; // [cluster] -> [dir path] TString TmpDir; bool KeepTempTables; diff --git a/ydb/library/yql/tests/s-expressions/suites/ManyInputTables/SplitTableRange.yql b/ydb/library/yql/tests/s-expressions/suites/ManyInputTables/SplitTableRange.yql index 98c32578b8..8793b592e7 100644 --- a/ydb/library/yql/tests/s-expressions/suites/ManyInputTables/SplitTableRange.yql +++ b/ydb/library/yql/tests/s-expressions/suites/ManyInputTables/SplitTableRange.yql @@ -2,7 +2,7 @@ #comment (let mr_source (DataSource 'yt 'plato)) (let world (Configure! world (DataSource '"yt" '"$all") '"Attr" '"maxinputtables" '"2")) -(let x (Read! world mr_source (Key '('table (MrTableRange '""))) (Void) '())) +(let x (Read! world mr_source (Key '('table (MrTableRange '"" (lambda '($i) (And (>= $i (String '"Input1")) (<= $i (String '"Input~")))) '""))) (Void) '())) (let world (Left! x)) (let tables (Right! x)) (let tables (Sort tables '((Bool 'true) (Bool 'true)) (lambda '(item) '((Member item 'key) (Member item 'subkey))))) diff --git a/ydb/library/yql/tests/s-expressions/suites/Udf/RecordRemapRangeDiffUdf.yql b/ydb/library/yql/tests/s-expressions/suites/Udf/RecordRemapRangeDiffUdf.yql index 8da6e37c67..fad593cd98 100644 --- a/ydb/library/yql/tests/s-expressions/suites/Udf/RecordRemapRangeDiffUdf.yql +++ b/ydb/library/yql/tests/s-expressions/suites/Udf/RecordRemapRangeDiffUdf.yql @@ -2,7 +2,7 @@ #comment (let mr_source (DataSource 'yt 'plato)) (let x (Read! world mr_source -(Key '('table (MrTableRange '""))) +(Key '('table (MrTableRange '"" (lambda '($i) (And (>= $i (String '"Input1")) (<= $i (String '"Input~")))) '""))) '('key 'subkey 'value) '())) (let world (Left! x)) (let table (Right! x)) diff --git a/ydb/library/yql/tests/s-expressions/suites/Udf/RecordRemapRangeDiffUdfPartial.yql b/ydb/library/yql/tests/s-expressions/suites/Udf/RecordRemapRangeDiffUdfPartial.yql index 22104d3cdf..d86392d08a 100644 --- a/ydb/library/yql/tests/s-expressions/suites/Udf/RecordRemapRangeDiffUdfPartial.yql +++ b/ydb/library/yql/tests/s-expressions/suites/Udf/RecordRemapRangeDiffUdfPartial.yql @@ -2,7 +2,7 @@ #comment (let mr_source (DataSource 'yt 'plato)) (let x (Read! world mr_source -(Key '('table (MrTableRange '""))) +(Key '('table (MrTableRange '"" (lambda '($i) (And (>= $i (String '"Input1")) (<= $i (String '"Input~")))) '""))) '('key 'subkey 'value) '())) (let world (Left! x)) (let table (Right! x)) diff --git a/ydb/library/yql/tests/s-expressions/suites/Udf/RecordRemapRangeSameUdf.yql b/ydb/library/yql/tests/s-expressions/suites/Udf/RecordRemapRangeSameUdf.yql index 8da6e37c67..fad593cd98 100644 --- a/ydb/library/yql/tests/s-expressions/suites/Udf/RecordRemapRangeSameUdf.yql +++ b/ydb/library/yql/tests/s-expressions/suites/Udf/RecordRemapRangeSameUdf.yql @@ -2,7 +2,7 @@ #comment (let mr_source (DataSource 'yt 'plato)) (let x (Read! world mr_source -(Key '('table (MrTableRange '""))) +(Key '('table (MrTableRange '"" (lambda '($i) (And (>= $i (String '"Input1")) (<= $i (String '"Input~")))) '""))) '('key 'subkey 'value) '())) (let world (Left! x)) (let table (Right! x)) diff --git a/ydb/library/yql/tests/s-expressions/suites/Udf/RecordRemapWeakRange.yql b/ydb/library/yql/tests/s-expressions/suites/Udf/RecordRemapWeakRange.yql index c3a1529c3d..09a346ac6d 100644 --- a/ydb/library/yql/tests/s-expressions/suites/Udf/RecordRemapWeakRange.yql +++ b/ydb/library/yql/tests/s-expressions/suites/Udf/RecordRemapWeakRange.yql @@ -1,7 +1,7 @@ ( #comment (let mr_source (DataSource 'yt 'plato)) -(let x (Read! world mr_source (Key '('table (MrTableRange '""))) '('key) '())) +(let x (Read! world mr_source (Key '('table (MrTableRange '"" (lambda '($i) (And (>= $i (String '"Input")) (<= $i (String '"Input~")))) '""))) '('key) '())) (let world (Left! x)) (let table1 (Right! x)) (let res_sink (DataSink 'result)) diff --git a/ydb/library/yql/tests/sql/dq_file/part3/canondata/result.json b/ydb/library/yql/tests/sql/dq_file/part3/canondata/result.json index 25bf209151..29bfb3a7ab 100644 --- a/ydb/library/yql/tests/sql/dq_file/part3/canondata/result.json +++ b/ydb/library/yql/tests/sql/dq_file/part3/canondata/result.json @@ -95,9 +95,9 @@ ], "test.test[action-subquery_merge_nested_subquery-default.txt-Debug]": [ { - "checksum": "87afe2e82ba47dd25434edfbf30701f9", - "size": 467, - "uri": "https://{canondata_backend}/1899731/5f48750839c300c592c921895adce61b6bdd10c7/resource.tar.gz#test.test_action-subquery_merge_nested_subquery-default.txt-Debug_/opt.yql_patched" + "checksum": "286067a89e2bb5a177c709e7eefa85de", + "size": 621, + "uri": "https://{canondata_backend}/1881367/d9d29884f2b57bc77cedad27e7b24a2adc89e30d/resource.tar.gz#test.test_action-subquery_merge_nested_subquery-default.txt-Debug_/opt.yql_patched" } ], "test.test[action-subquery_merge_nested_subquery-default.txt-Plan]": [ diff --git a/ydb/library/yql/tests/sql/hybrid_file/part2/canondata/result.json b/ydb/library/yql/tests/sql/hybrid_file/part2/canondata/result.json index 0ce7270cf8..5cb7c959b9 100644 --- a/ydb/library/yql/tests/sql/hybrid_file/part2/canondata/result.json +++ b/ydb/library/yql/tests/sql/hybrid_file/part2/canondata/result.json @@ -57,9 +57,9 @@ ], "test.test[action-subquery_merge_nested_subquery-default.txt-Debug]": [ { - "checksum": "d32a3afdd03b6999f6e278e75caf9d06", - "size": 466, - "uri": "https://{canondata_backend}/1936842/51593b2a750dbb036388d012a30fa937edaab5f0/resource.tar.gz#test.test_action-subquery_merge_nested_subquery-default.txt-Debug_/opt.yql_patched" + "checksum": "9b9bd9b5c87918ed2451269a9938d2c4", + "size": 620, + "uri": "https://{canondata_backend}/1942278/8b999765d5b3ee30766f09cb37b290aab0244d6e/resource.tar.gz#test.test_action-subquery_merge_nested_subquery-default.txt-Debug_/opt.yql_patched" } ], "test.test[action-subquery_merge_nested_subquery-default.txt-Plan]": [ diff --git a/ydb/library/yql/tests/sql/sql2yql/canondata/result.json b/ydb/library/yql/tests/sql/sql2yql/canondata/result.json index 26a3ca6a4d..ae590614b4 100644 --- a/ydb/library/yql/tests/sql/sql2yql/canondata/result.json +++ b/ydb/library/yql/tests/sql/sql2yql/canondata/result.json @@ -456,9 +456,9 @@ ], "test_sql2yql.test[action-insert_each_from_folder]": [ { - "checksum": "dfda31784ae59d009f003ee44ca0f4a2", - "size": 3140, - "uri": "https://{canondata_backend}/1937027/973c239492ba32946806ddc66cf0af4b38c06ae8/resource.tar.gz#test_sql2yql.test_action-insert_each_from_folder_/sql.yql" + "checksum": "7c5460adea49ac910a8da3e167784434", + "size": 3446, + "uri": "https://{canondata_backend}/1920236/49d24809be769996fac3586d53f86d8c72bdeb78/resource.tar.gz#test_sql2yql.test_action-insert_each_from_folder_/sql.yql" } ], "test_sql2yql.test[action-lambda_arg_count]": [ @@ -694,9 +694,9 @@ ], "test_sql2yql.test[action-subquery_merge_nested_world]": [ { - "checksum": "0eba6f388da133f0f09d55d25be70567", - "size": 3226, - "uri": "https://{canondata_backend}/1784117/d56ae82ad9d30397a41490647be1bd2124718f98/resource.tar.gz#test_sql2yql.test_action-subquery_merge_nested_world_/sql.yql" + "checksum": "54c3a4cecc1ea4b1bf1fc7ed4b67d9a9", + "size": 3613, + "uri": "https://{canondata_backend}/1920236/49d24809be769996fac3586d53f86d8c72bdeb78/resource.tar.gz#test_sql2yql.test_action-subquery_merge_nested_world_/sql.yql" } ], "test_sql2yql.test[action-subquery_opt_args]": [ @@ -729,9 +729,9 @@ ], "test_sql2yql.test[action-table_content_before_from_folder]": [ { - "checksum": "e4fed412d365d928d7b9d1c18af8144d", - "size": 3452, - "uri": "https://{canondata_backend}/1784117/d56ae82ad9d30397a41490647be1bd2124718f98/resource.tar.gz#test_sql2yql.test_action-table_content_before_from_folder_/sql.yql" + "checksum": "74563a298fcb8f48e9090ad7d5d29add", + "size": 3831, + "uri": "https://{canondata_backend}/1920236/49d24809be769996fac3586d53f86d8c72bdeb78/resource.tar.gz#test_sql2yql.test_action-table_content_before_from_folder_/sql.yql" } ], "test_sql2yql.test[agg_apply-avg_const_interval]": [ @@ -19447,9 +19447,9 @@ ], "test_sql_format.test[action-insert_each_from_folder]": [ { - "checksum": "42fbd294c401e7a47b4319533de45aee", - "size": 366, - "uri": "https://{canondata_backend}/1880306/64654158d6bfb1289c66c626a8162239289559d0/resource.tar.gz#test_sql_format.test_action-insert_each_from_folder_/formatted.sql" + "checksum": "c51ad51888f8e943d2edc55f296bf202", + "size": 389, + "uri": "https://{canondata_backend}/1925821/6c2f883a1c33f02b8bcef287229b2b73dd762cff/resource.tar.gz#test_sql_format.test_action-insert_each_from_folder_/formatted.sql" } ], "test_sql_format.test[action-lambda_arg_count]": [ @@ -19685,9 +19685,9 @@ ], "test_sql_format.test[action-subquery_merge_nested_world]": [ { - "checksum": "02cdaf09ad3c7c45b643f6e37cd099f1", - "size": 292, - "uri": "https://{canondata_backend}/1880306/64654158d6bfb1289c66c626a8162239289559d0/resource.tar.gz#test_sql_format.test_action-subquery_merge_nested_world_/formatted.sql" + "checksum": "8c0ad08eb1ba55aa5f7c977d6029a8c7", + "size": 325, + "uri": "https://{canondata_backend}/1925821/6c2f883a1c33f02b8bcef287229b2b73dd762cff/resource.tar.gz#test_sql_format.test_action-subquery_merge_nested_world_/formatted.sql" } ], "test_sql_format.test[action-subquery_opt_args]": [ @@ -19720,9 +19720,9 @@ ], "test_sql_format.test[action-table_content_before_from_folder]": [ { - "checksum": "55901f202fac5c4226ff35047872c116", - "size": 258, - "uri": "https://{canondata_backend}/1880306/64654158d6bfb1289c66c626a8162239289559d0/resource.tar.gz#test_sql_format.test_action-table_content_before_from_folder_/formatted.sql" + "checksum": "45dce176f9281a8754c8da9514085c3f", + "size": 287, + "uri": "https://{canondata_backend}/1925821/6c2f883a1c33f02b8bcef287229b2b73dd762cff/resource.tar.gz#test_sql_format.test_action-table_content_before_from_folder_/formatted.sql" } ], "test_sql_format.test[agg_apply-avg_const_interval]": [ diff --git a/ydb/library/yql/tests/sql/suites/action/insert_each_from_folder.sql b/ydb/library/yql/tests/sql/suites/action/insert_each_from_folder.sql index 4fdbd460d6..dc06200679 100644 --- a/ydb/library/yql/tests/sql/suites/action/insert_each_from_folder.sql +++ b/ydb/library/yql/tests/sql/suites/action/insert_each_from_folder.sql @@ -6,7 +6,7 @@ use plato; $list = ( select aggregate_list(Path) from ( select Path from folder("") - where Type = "table" + where Type = "table" and Path like "Input%" order by Path desc limit 30 ) diff --git a/ydb/library/yql/tests/sql/suites/action/subquery_merge_nested_world.sql b/ydb/library/yql/tests/sql/suites/action/subquery_merge_nested_world.sql index e31e1ffbac..67f0683145 100644 --- a/ydb/library/yql/tests/sql/suites/action/subquery_merge_nested_world.sql +++ b/ydb/library/yql/tests/sql/suites/action/subquery_merge_nested_world.sql @@ -3,7 +3,7 @@ use plato; DEFINE SUBQUERY $s($_i) AS - $t = SELECT AGGREGATE_LIST(Path) FROM FOLDER(''); + $t = SELECT AGGREGATE_LIST(Path) FROM FOLDER('') WHERE Path LIKE "Input%"; SELECT * FROM EACH($t); diff --git a/ydb/library/yql/tests/sql/suites/action/table_content_before_from_folder.sql b/ydb/library/yql/tests/sql/suites/action/table_content_before_from_folder.sql index 5a1da378af..bd6ee020d4 100644 --- a/ydb/library/yql/tests/sql/suites/action/table_content_before_from_folder.sql +++ b/ydb/library/yql/tests/sql/suites/action/table_content_before_from_folder.sql @@ -3,5 +3,5 @@ use plato; pragma yt.EvaluationTableSizeLimit="1"; select * from Input limit 1; -$tables = (select aggregate_list(Path) as dates from folder("")); +$tables = (select aggregate_list(Path) as dates from folder("") where Path like "Input%"); select count(*) from each($tables); diff --git a/ydb/library/yql/tests/sql/yt_native_file/part3/canondata/result.json b/ydb/library/yql/tests/sql/yt_native_file/part3/canondata/result.json index c8d04b7172..c94a3a82ee 100644 --- a/ydb/library/yql/tests/sql/yt_native_file/part3/canondata/result.json +++ b/ydb/library/yql/tests/sql/yt_native_file/part3/canondata/result.json @@ -78,9 +78,9 @@ ], "test.test[action-subquery_merge_nested_subquery-default.txt-Debug]": [ { - "checksum": "1dc78c0483e5fa3351b051e8e9de8679", - "size": 399, - "uri": "https://{canondata_backend}/1924537/3d2705efadc402a8a106ede4d6132e0dbc7b7516/resource.tar.gz#test.test_action-subquery_merge_nested_subquery-default.txt-Debug_/opt.yql" + "checksum": "77b06488833b67e7d27ef6221c3418e5", + "size": 553, + "uri": "https://{canondata_backend}/1920236/df296391156851095e326102a665e0179365e8d4/resource.tar.gz#test.test_action-subquery_merge_nested_subquery-default.txt-Debug_/opt.yql" } ], "test.test[action-subquery_merge_nested_subquery-default.txt-Plan]": [ @@ -92,9 +92,9 @@ ], "test.test[action-subquery_merge_nested_subquery-default.txt-Results]": [ { - "checksum": "ab3678f99412b6d14f44db6b69d41e5f", - "size": 668, - "uri": "https://{canondata_backend}/1871182/5a4449eecae85d31c6f70eca057fe425dfafe11b/resource.tar.gz#test.test_action-subquery_merge_nested_subquery-default.txt-Results_/results.txt" + "checksum": "6a84f2308a09ca7b3f129ecebf10f516", + "size": 747, + "uri": "https://{canondata_backend}/1920236/df296391156851095e326102a665e0179365e8d4/resource.tar.gz#test.test_action-subquery_merge_nested_subquery-default.txt-Results_/results.txt" } ], "test.test[agg_phases-min_by-default.txt-Debug]": [ diff --git a/ydb/library/yql/tools/yqlrun/yqlrun.cpp b/ydb/library/yql/tools/yqlrun/yqlrun.cpp index 4c92de66b6..dcc03d1d95 100644 --- a/ydb/library/yql/tools/yqlrun/yqlrun.cpp +++ b/ydb/library/yql/tools/yqlrun/yqlrun.cpp @@ -1,5 +1,6 @@ #include "gateway_spec.h" +#include <filesystem> #include <ydb/library/yql/tools/yqlrun/http/yql_server.h> #include <ydb/library/yql/providers/yt/gateway/file/yql_yt_file.h> @@ -380,7 +381,9 @@ int Main(int argc, const char *argv[]) TOpts opts = TOpts::Default(); TString programFile; TVector<TString> tablesMappingList; + TVector<TString> tablesDirMappingList; THashMap<TString, TString> tablesMapping; + THashMap<TString, TString> tablesDirMapping; TVector<TString> filesMappingList; TUserDataTable filesMapping; TVector<TString> urlsMappingList; @@ -413,6 +416,7 @@ int Main(int argc, const char *argv[]) opts.AddLongOption('s', "sql", "program is SQL query").NoArgument(); opts.AddLongOption("pg", "program has PG syntax").NoArgument(); opts.AddLongOption('t', "table", "table@file").AppendTo(&tablesMappingList); + opts.AddLongOption("tables-dir", "cluster@dir").AppendTo(&tablesDirMappingList); opts.AddLongOption('C', "cluster", "set cluster to service mapping").RequiredArgument("name@service").Handler(new TStoreMappingFunctor(&clusterMapping)); opts.AddLongOption("ndebug", "should be at first argument, do not show debug info in error output").NoArgument(); opts.AddLongOption("parse-only", "exit after program has been parsed").NoArgument(); @@ -502,6 +506,23 @@ int Main(int argc, const char *argv[]) tablesMapping[tableName] = filePath; } + for (auto& s : tablesDirMappingList) { + TStringBuf clusterName, dirPath; + TStringBuf(s).Split('@', clusterName, dirPath); + if (clusterName.empty() || dirPath.empty()) { + Cerr << "Incorrect table directory mapping, expected form cluster@dir, e.g. yt.plato@/tmp/tables" << Endl; + return 1; + } + tablesDirMapping[clusterName] = dirPath; + for (const auto& entry : std::filesystem::recursive_directory_iterator(std::string(dirPath))) { + if (entry.is_regular_file() && entry.path().has_extension() && entry.path().extension() == ".txt") { + auto tableName = TString(clusterName) + '.' + std::filesystem::relative(entry.path(), std::string(dirPath)); + tableName = tableName.substr(0, tableName.Size() - 4); // remove .txt extension + tablesMapping[tableName] = entry.path().string(); + } + } + } + if (hasValidate) { for (auto& s : filesMappingList) { TStringBuf fileName, filePath; @@ -644,7 +665,7 @@ int Main(int argc, const char *argv[]) bool emulateOutputForMultirun = false; if (hasValidate) { if (gatewayTypes.contains(YtProviderName) || res.Has("opt-collision")) { - auto yqlNativeServices = NFile::TYtFileServices::Make(funcRegistry.Get(), tablesMapping, fileStorage, tmpDir, res.Has("keep-temp")); + auto yqlNativeServices = NFile::TYtFileServices::Make(funcRegistry.Get(), tablesMapping, fileStorage, tmpDir, res.Has("keep-temp"), tablesDirMapping); auto ytNativeGateway = CreateYtFileGateway(yqlNativeServices, &emulateOutputForMultirun); dataProvidersInit.push_back(GetYtNativeDataProviderInitializer(ytNativeGateway)); } |