Skip to content

Commit aad4438

Browse files
mbasmanovafacebook-github-bot
mbasmanova
authored and
facebook-github-bot
committed
misc: Extract TestIndexTable::create helper method (facebookincubator#13852)
Summary: Pull Request resolved: facebookincubator#13852 Allow to reuse this logic in multiple tests. Reviewed By: xiaoxmeng Differential Revision: D77160446 fbshipit-source-id: 4af4b93e65a832c6aca7d897166d4a74b216715d
1 parent 1769ca7 commit aad4438

File tree

3 files changed

+116
-89
lines changed

3 files changed

+116
-89
lines changed

‎velox/exec/tests/IndexLookupJoinTest.cpp

Lines changed: 35 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -113,74 +113,6 @@ class IndexLookupJoinTest : public IndexLookupJoinTestBase,
113113
ASSERT_EQ(plan->toString(true, true), copy->toString(true, true));
114114
}
115115

116-
// Create index table with the given key and value inputs.
117-
std::shared_ptr<TestIndexTable> createIndexTable(
118-
int numEqualJoinKeys,
119-
const RowVectorPtr& keyData,
120-
const RowVectorPtr& valueData) {
121-
const auto keyType =
122-
std::dynamic_pointer_cast<const RowType>(keyData->type());
123-
VELOX_CHECK_GE(keyType->size(), 1);
124-
VELOX_CHECK_GE(keyType->size(), numEqualJoinKeys);
125-
auto valueType =
126-
std::dynamic_pointer_cast<const RowType>(valueData->type());
127-
VELOX_CHECK_GE(valueType->size(), 1);
128-
const auto numRows = keyData->size();
129-
VELOX_CHECK_EQ(numRows, valueData->size());
130-
131-
std::vector<std::unique_ptr<VectorHasher>> hashers;
132-
hashers.reserve(numEqualJoinKeys);
133-
std::vector<VectorPtr> keyVectors;
134-
keyVectors.reserve(numEqualJoinKeys);
135-
for (auto i = 0; i < numEqualJoinKeys; ++i) {
136-
hashers.push_back(std::make_unique<VectorHasher>(keyType->childAt(i), i));
137-
keyVectors.push_back(keyData->childAt(i));
138-
}
139-
140-
std::vector<TypePtr> dependentTypes;
141-
std::vector<VectorPtr> dependentVectors;
142-
for (int i = numEqualJoinKeys; i < keyType->size(); ++i) {
143-
dependentTypes.push_back(keyType->childAt(i));
144-
dependentVectors.push_back(keyData->childAt(i));
145-
}
146-
for (int i = 0; i < valueType->size(); ++i) {
147-
dependentTypes.push_back(valueType->childAt(i));
148-
dependentVectors.push_back(valueData->childAt(i));
149-
}
150-
151-
// Create the table.
152-
auto table = HashTable<false>::createForJoin(
153-
std::move(hashers),
154-
/*dependentTypes=*/dependentTypes,
155-
/*allowDuplicates=*/true,
156-
/*hasProbedFlag=*/false,
157-
/*minTableSizeForParallelJoinBuild=*/1,
158-
pool_.get());
159-
160-
// Insert data into the row container.
161-
auto* rowContainer = table->rows();
162-
std::vector<DecodedVector> decodedVectors;
163-
for (auto& vector : keyData->children()) {
164-
decodedVectors.emplace_back(*vector);
165-
}
166-
for (auto& vector : valueData->children()) {
167-
decodedVectors.emplace_back(*vector);
168-
}
169-
170-
for (auto row = 0; row < numRows; ++row) {
171-
auto* newRow = rowContainer->newRow();
172-
173-
for (auto col = 0; col < decodedVectors.size(); ++col) {
174-
rowContainer->store(decodedVectors[col], row, newRow, col);
175-
}
176-
}
177-
178-
// Build the table index.
179-
table->prepareJoinTable({}, BaseHashTable::kNoSpillInputStartPartitionBit);
180-
return std::make_shared<TestIndexTable>(
181-
std::move(keyType), std::move(valueType), std::move(table));
182-
}
183-
184116
// Makes index table handle with the specified index table and async lookup
185117
// flag.
186118
std::shared_ptr<TestIndexTableHandle> makeIndexTableHandle(
@@ -795,8 +727,11 @@ TEST_P(IndexLookupJoinTest, equalJoin) {
795727
createDuckDbTable("t", probeVectors);
796728
createDuckDbTable("u", {tableData.tableData});
797729

798-
const auto indexTable = createIndexTable(
799-
/*numEqualJoinKeys=*/3, tableData.keyData, tableData.valueData);
730+
const auto indexTable = TestIndexTable::create(
731+
/*numEqualJoinKeys=*/3,
732+
tableData.keyData,
733+
tableData.valueData,
734+
*pool());
800735
const auto indexTableHandle =
801736
makeIndexTableHandle(indexTable, GetParam().asyncLookup);
802737
auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
@@ -1250,8 +1185,11 @@ TEST_P(IndexLookupJoinTest, betweenJoinCondition) {
12501185
createDuckDbTable("t", probeVectors);
12511186
createDuckDbTable("u", {tableData.tableData});
12521187

1253-
const auto indexTable = createIndexTable(
1254-
/*numEqualJoinKeys=*/2, tableData.keyData, tableData.valueData);
1188+
const auto indexTable = TestIndexTable::create(
1189+
/*numEqualJoinKeys=*/2,
1190+
tableData.keyData,
1191+
tableData.valueData,
1192+
*pool());
12551193
const auto indexTableHandle =
12561194
makeIndexTableHandle(indexTable, GetParam().asyncLookup);
12571195
auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
@@ -1571,8 +1509,11 @@ TEST_P(IndexLookupJoinTest, inJoinCondition) {
15711509
createDuckDbTable("t", probeVectors);
15721510
createDuckDbTable("u", {tableData.tableData});
15731511

1574-
const auto indexTable = createIndexTable(
1575-
/*numEqualJoinKeys=*/2, tableData.keyData, tableData.valueData);
1512+
const auto indexTable = TestIndexTable::create(
1513+
/*numEqualJoinKeys=*/2,
1514+
tableData.keyData,
1515+
tableData.valueData,
1516+
*pool());
15761517
const auto indexTableHandle =
15771518
makeIndexTableHandle(indexTable, GetParam().asyncLookup);
15781519
auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
@@ -1622,8 +1563,8 @@ DEBUG_ONLY_TEST_P(IndexLookupJoinTest, connectorError) {
16221563
}
16231564
}));
16241565

1625-
const auto indexTable = createIndexTable(
1626-
/*numEqualJoinKeys=*/3, tableData.keyData, tableData.valueData);
1566+
const auto indexTable = TestIndexTable::create(
1567+
/*numEqualJoinKeys=*/3, tableData.keyData, tableData.valueData, *pool());
16271568
const auto indexTableHandle =
16281569
makeIndexTableHandle(indexTable, GetParam().asyncLookup);
16291570
auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
@@ -1691,8 +1632,8 @@ DEBUG_ONLY_TEST_P(IndexLookupJoinTest, prefetch) {
16911632
asyncLookupWait.await([&] { return !asyncLookupWaitFlag.load(); });
16921633
}));
16931634

1694-
const auto indexTable = createIndexTable(
1695-
/*numEqualJoinKeys=*/3, tableData.keyData, tableData.valueData);
1635+
const auto indexTable = TestIndexTable::create(
1636+
/*numEqualJoinKeys=*/3, tableData.keyData, tableData.valueData, *pool());
16961637
const auto indexTableHandle =
16971638
makeIndexTableHandle(indexTable, GetParam().asyncLookup);
16981639
auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
@@ -1790,8 +1731,11 @@ TEST_P(IndexLookupJoinTest, outputBatchSizeWithInnerJoin) {
17901731
createDuckDbTable("t", probeVectors);
17911732
createDuckDbTable("u", {tableData.tableData});
17921733

1793-
const auto indexTable = createIndexTable(
1794-
/*numEqualJoinKeys=*/3, tableData.keyData, tableData.valueData);
1734+
const auto indexTable = TestIndexTable::create(
1735+
/*numEqualJoinKeys=*/3,
1736+
tableData.keyData,
1737+
tableData.valueData,
1738+
*pool());
17951739
const auto indexTableHandle =
17961740
makeIndexTableHandle(indexTable, GetParam().asyncLookup);
17971741
auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
@@ -1895,8 +1839,11 @@ TEST_P(IndexLookupJoinTest, outputBatchSizeWithLeftJoin) {
18951839
createDuckDbTable("t", probeVectors);
18961840
createDuckDbTable("u", {tableData.tableData});
18971841

1898-
const auto indexTable = createIndexTable(
1899-
/*numEqualJoinKeys=*/3, tableData.keyData, tableData.valueData);
1842+
const auto indexTable = TestIndexTable::create(
1843+
/*numEqualJoinKeys=*/3,
1844+
tableData.keyData,
1845+
tableData.valueData,
1846+
*pool());
19001847
const auto indexTableHandle =
19011848
makeIndexTableHandle(indexTable, GetParam().asyncLookup);
19021849
auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
@@ -1967,8 +1914,8 @@ DEBUG_ONLY_TEST_P(IndexLookupJoinTest, runtimeStats) {
19671914
std::this_thread::sleep_for(std::chrono::milliseconds(100)); // NOLINT
19681915
}));
19691916

1970-
const auto indexTable = createIndexTable(
1971-
/*numEqualJoinKeys=*/3, tableData.keyData, tableData.valueData);
1917+
const auto indexTable = TestIndexTable::create(
1918+
/*numEqualJoinKeys=*/3, tableData.keyData, tableData.valueData, *pool());
19721919
const auto indexTableHandle =
19731920
makeIndexTableHandle(indexTable, GetParam().asyncLookup);
19741921
auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
@@ -2057,8 +2004,8 @@ TEST_P(IndexLookupJoinTest, barrier) {
20572004
createDuckDbTable("t", probeVectors);
20582005
createDuckDbTable("u", {tableData.tableData});
20592006

2060-
const auto indexTable = createIndexTable(
2061-
/*numEqualJoinKeys=*/3, tableData.keyData, tableData.valueData);
2007+
const auto indexTable = TestIndexTable::create(
2008+
/*numEqualJoinKeys=*/3, tableData.keyData, tableData.valueData, *pool());
20622009
const auto indexTableHandle =
20632010
makeIndexTableHandle(indexTable, GetParam().asyncLookup);
20642011
auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
@@ -2128,8 +2075,8 @@ TEST_P(IndexLookupJoinTest, joinFuzzer) {
21282075
createDuckDbTable("t", probeVectors);
21292076
createDuckDbTable("u", {tableData.tableData});
21302077

2131-
const auto indexTable = createIndexTable(
2132-
/*numEqualJoinKeys=*/1, tableData.keyData, tableData.valueData);
2078+
const auto indexTable = TestIndexTable::create(
2079+
/*numEqualJoinKeys=*/1, tableData.keyData, tableData.valueData, *pool());
21332080
const auto indexTableHandle =
21342081
makeIndexTableHandle(indexTable, GetParam().asyncLookup);
21352082
auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();

‎velox/exec/tests/utils/TestIndexStorageConnector.cpp

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,80 @@
2525

2626
using facebook::velox::common::testutil::TestValue;
2727
namespace facebook::velox::exec::test {
28+
29+
// static
30+
std::shared_ptr<TestIndexTable> TestIndexTable::create(
31+
size_t numEqualJoinKeys,
32+
const RowVectorPtr& keyData,
33+
const RowVectorPtr& valueData,
34+
velox::memory::MemoryPool& pool) {
35+
VELOX_CHECK_GE(numEqualJoinKeys, 1);
36+
VELOX_CHECK_NOT_NULL(keyData);
37+
VELOX_CHECK_NOT_NULL(valueData);
38+
39+
auto keyType = asRowType(keyData->type());
40+
VELOX_CHECK_GE(keyType->size(), numEqualJoinKeys);
41+
42+
auto valueType = asRowType(valueData->type());
43+
VELOX_CHECK_GE(valueType->size(), 1);
44+
45+
const auto numRows = keyData->size();
46+
VELOX_CHECK_EQ(numRows, valueData->size());
47+
48+
std::vector<std::unique_ptr<VectorHasher>> hashers;
49+
hashers.reserve(numEqualJoinKeys);
50+
std::vector<VectorPtr> keyVectors;
51+
keyVectors.reserve(numEqualJoinKeys);
52+
for (auto i = 0; i < numEqualJoinKeys; ++i) {
53+
hashers.push_back(std::make_unique<VectorHasher>(keyType->childAt(i), i));
54+
keyVectors.push_back(keyData->childAt(i));
55+
}
56+
57+
std::vector<TypePtr> dependentTypes;
58+
std::vector<VectorPtr> dependentVectors;
59+
for (auto i = numEqualJoinKeys; i < keyType->size(); ++i) {
60+
dependentTypes.push_back(keyType->childAt(i));
61+
dependentVectors.push_back(keyData->childAt(i));
62+
}
63+
64+
for (auto i = 0; i < valueType->size(); ++i) {
65+
dependentTypes.push_back(valueType->childAt(i));
66+
dependentVectors.push_back(valueData->childAt(i));
67+
}
68+
69+
// Create the table.
70+
auto table = HashTable<false>::createForJoin(
71+
std::move(hashers),
72+
/*dependentTypes=*/dependentTypes,
73+
/*allowDuplicates=*/true,
74+
/*hasProbedFlag=*/false,
75+
/*minTableSizeForParallelJoinBuild=*/1,
76+
&pool);
77+
78+
// Insert data into the row container.
79+
auto* rowContainer = table->rows();
80+
std::vector<DecodedVector> decodedVectors;
81+
for (auto& vector : keyData->children()) {
82+
decodedVectors.emplace_back(*vector);
83+
}
84+
for (auto& vector : valueData->children()) {
85+
decodedVectors.emplace_back(*vector);
86+
}
87+
88+
for (auto row = 0; row < numRows; ++row) {
89+
auto* newRow = rowContainer->newRow();
90+
91+
for (auto col = 0; col < decodedVectors.size(); ++col) {
92+
rowContainer->store(decodedVectors[col], row, newRow, col);
93+
}
94+
}
95+
96+
// Build the table index.
97+
table->prepareJoinTable({}, BaseHashTable::kNoSpillInputStartPartitionBit);
98+
return std::make_shared<TestIndexTable>(
99+
std::move(keyType), std::move(valueType), std::move(table));
100+
}
101+
28102
namespace {
29103
core::TypedExprPtr toJoinConditionExpr(
30104
const std::vector<std::shared_ptr<core::IndexLookupCondition>>&

‎velox/exec/tests/utils/TestIndexStorageConnector.h

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
#pragma once
1717

1818
#include "velox/exec/HashTable.h"
19-
#include "velox/exec/OperatorUtils.h"
2019
#include "velox/type/Type.h"
2120

2221
namespace facebook::velox::exec::test {
@@ -38,6 +37,13 @@ struct TestIndexTable {
3837
: keyType(std::move(_keyType)),
3938
dataType(std::move(_dataType)),
4039
table(std::move(_table)) {}
40+
41+
// Create index table with the given key and value inputs.
42+
static std::shared_ptr<TestIndexTable> create(
43+
size_t numEqualJoinKeys,
44+
const RowVectorPtr& keyData,
45+
const RowVectorPtr& valueData,
46+
velox::memory::MemoryPool& pool);
4147
};
4248

4349
// The index table handle which provides the index table for index lookup.

0 commit comments

Comments
 (0)