Skip to content

Commit

Permalink
feat(fuzzer): Support multiple joins in the join node "toSql" methods…
Browse files Browse the repository at this point in the history 
… for reference query runners (facebookincubator#11801)

Summary:

Currently, the hash join and nested loop join "toSql" methods for all reference query runners only support a single join. This change extends it to support multiple joins, only needing the join node of the last join in the tree. It traverses up the tree and recursively builds the sql query.

Differential Revision: D66977480
  • Loading branch information
@facebook-github-bot
Daniel Hunte authored and facebook-github-bot committed Dec 27, 2024
1 parent 9e418b1 commit e222e6d
Show file tree
Hide file tree
Showing 8 changed files with 420 additions and 108 deletions.
2 changes: 2 additions & 0 deletions velox/core/PlanNode.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -324,6 +324,8 @@ class ValuesNode : public PlanNode {
const size_t repeatTimes_;
};

using ValuesNodePtr = std::shared_ptr<const ValuesNode>;

class ArrowStreamNode : public PlanNode {
public:
ArrowStreamNode(
Expand Down
160 changes: 113 additions & 47 deletions velox/exec/fuzzer/DuckQueryRunner.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,38 @@ DuckQueryRunner::aggregationFunctionDataSpecs() const {
return kAggregationFunctionDataSpecs;
}

std::string DuckQueryRunner::getTableName(
const core::ValuesNodePtr& valuesNode) {
return fmt::format("t_{}", valuesNode->id());
}

std::unordered_map<std::string, std::vector<velox::RowVectorPtr>>
DuckQueryRunner::getAllTablesAndNames(const core::PlanNodePtr& plan) {
std::unordered_map<std::string, std::vector<velox::RowVectorPtr>> result;
if (const auto valuesNode =
std::dynamic_pointer_cast<const core::ValuesNode>(plan)) {
result.insert({getTableName(valuesNode), valuesNode->values()});
} else {
for (const auto& source : plan->sources()) {
auto tablesAndNames = getAllTablesAndNames(source);
result.insert(tablesAndNames.begin(), tablesAndNames.end());
}
}
return result;
}

std::multiset<std::vector<velox::variant>> DuckQueryRunner::execute(
const std::string& sql,
const core::PlanNodePtr& plan) {
DuckDbQueryRunner queryRunner;
std::unordered_map<std::string, std::vector<RowVectorPtr>> inputMap =
getAllTablesAndNames(plan);
for (const auto& [tableName, input] : inputMap) {
queryRunner.createTable(tableName, input);
}
return queryRunner.execute(sql, plan->outputType());
}

std::multiset<std::vector<velox::variant>> DuckQueryRunner::execute(
const std::string& sql,
const std::vector<RowVectorPtr>& input,
Expand Down Expand Up @@ -341,38 +373,53 @@ std::optional<std::string> DuckQueryRunner::toSql(
return sql.str();
}

std::optional<std::string> DuckQueryRunner::toSql(
const std::shared_ptr<const core::HashJoinNode>& joinNode) {
const auto& joinKeysToSql = [](auto keys) {
std::stringstream out;
for (auto i = 0; i < keys.size(); ++i) {
if (i > 0) {
out << ", ";
}
out << keys[i]->name();
static const std::string joinKeysToSql(
const std::vector<core::FieldAccessTypedExprPtr>& keys) {
std::stringstream out;
for (auto i = 0; i < keys.size(); ++i) {
if (i > 0) {
out << ", ";
}
return out.str();
};
out << keys[i]->name();
}
return out.str();
}

const auto filterToSql = [](core::TypedExprPtr filter) {
auto call = std::dynamic_pointer_cast<const core::CallTypedExpr>(filter);
return toCallSql(call);
};
static std::string filterToSql(const core::TypedExprPtr& filter) {
auto call = std::dynamic_pointer_cast<const core::CallTypedExpr>(filter);
return toCallSql(call);
}

const auto& joinConditionAsSql = [&](auto joinNode) {
std::stringstream out;
for (auto i = 0; i < joinNode->leftKeys().size(); ++i) {
if (i > 0) {
out << " AND ";
}
out << joinNode->leftKeys()[i]->name() << " = "
<< joinNode->rightKeys()[i]->name();
static std::string joinConditionAsSql(const core::AbstractJoinNode& joinNode) {
std::stringstream out;
for (auto i = 0; i < joinNode.leftKeys().size(); ++i) {
if (i > 0) {
out << " AND ";
}
if (joinNode->filter()) {
out << " AND " << filterToSql(joinNode->filter());
out << joinNode.leftKeys()[i]->name() << " = "
<< joinNode.rightKeys()[i]->name();
}
if (joinNode.filter()) {
if (!joinNode.leftKeys().empty()) {
out << " AND ";
}
return out.str();
};
out << filterToSql(joinNode.filter());
}
return out.str();
};

std::optional<std::string> DuckQueryRunner::toSql(
const std::shared_ptr<const core::HashJoinNode>& joinNode) {
std::string probeTableName;
std::string buildTableName;
if (const auto probeSubQuery = toSql(joinNode->sources()[0]),
buildSubQuery = toSql(joinNode->sources()[1]);
probeSubQuery && buildSubQuery) {
probeTableName = fmt::format("({})", *probeSubQuery);
buildTableName = fmt::format("({})", *buildSubQuery);
} else {
return std::nullopt;
}

const auto& outputNames = joinNode->outputType()->names();

Expand All @@ -386,24 +433,27 @@ std::optional<std::string> DuckQueryRunner::toSql(

switch (joinNode->joinType()) {
case core::JoinType::kInner:
sql << " FROM t INNER JOIN u ON " << joinConditionAsSql(joinNode);
sql << " FROM " << probeTableName << " INNER JOIN " << buildTableName
<< " ON " << joinConditionAsSql(*joinNode);
break;
case core::JoinType::kLeft:
sql << " FROM t LEFT JOIN u ON " << joinConditionAsSql(joinNode);
sql << " FROM " << probeTableName << " LEFT JOIN " << buildTableName
<< " ON " << joinConditionAsSql(*joinNode);
break;
case core::JoinType::kFull:
sql << " FROM t FULL OUTER JOIN u ON " << joinConditionAsSql(joinNode);
sql << " FROM " << probeTableName << " FULL OUTER JOIN " << buildTableName
<< " ON " << joinConditionAsSql(*joinNode);
break;
case core::JoinType::kLeftSemiFilter:
// Multiple columns returned by a scalar subquery is not supported in
// DuckDB. A scalar subquery expression is a subquery that returns one
// Presto. A scalar subquery expression is a subquery that returns one
// result row from exactly one column for every input row.
if (joinNode->leftKeys().size() > 1) {
return std::nullopt;
}
sql << " FROM t WHERE " << joinKeysToSql(joinNode->leftKeys())
<< " IN (SELECT " << joinKeysToSql(joinNode->rightKeys())
<< " FROM u";
sql << " FROM " << probeTableName << " WHERE "
<< joinKeysToSql(joinNode->leftKeys()) << " IN (SELECT "
<< joinKeysToSql(joinNode->rightKeys()) << " FROM " << buildTableName;
if (joinNode->filter()) {
sql << " WHERE " << filterToSql(joinNode->filter());
}
Expand All @@ -412,29 +462,31 @@ std::optional<std::string> DuckQueryRunner::toSql(
case core::JoinType::kLeftSemiProject:
if (joinNode->isNullAware()) {
sql << ", " << joinKeysToSql(joinNode->leftKeys()) << " IN (SELECT "
<< joinKeysToSql(joinNode->rightKeys()) << " FROM u";
<< joinKeysToSql(joinNode->rightKeys()) << " FROM "
<< buildTableName;
if (joinNode->filter()) {
sql << " WHERE " << filterToSql(joinNode->filter());
}
sql << ") FROM t";
sql << ") FROM " << probeTableName;
} else {
sql << ", EXISTS (SELECT * FROM u WHERE "
<< joinConditionAsSql(joinNode);
sql << ") FROM t";
sql << ", EXISTS (SELECT * FROM " << buildTableName << " WHERE "
<< joinConditionAsSql(*joinNode);
sql << ") FROM " << probeTableName;
}
break;
case core::JoinType::kAnti:
if (joinNode->isNullAware()) {
sql << " FROM t WHERE " << joinKeysToSql(joinNode->leftKeys())
<< " NOT IN (SELECT " << joinKeysToSql(joinNode->rightKeys())
<< " FROM u";
sql << " FROM " << probeTableName << " WHERE "
<< joinKeysToSql(joinNode->leftKeys()) << " NOT IN (SELECT "
<< joinKeysToSql(joinNode->rightKeys()) << " FROM "
<< buildTableName;
if (joinNode->filter()) {
sql << " WHERE " << filterToSql(joinNode->filter());
}
sql << ")";
} else {
sql << " FROM t WHERE NOT EXISTS (SELECT * FROM u WHERE "
<< joinConditionAsSql(joinNode);
sql << " FROM " << probeTableName << " WHERE NOT EXISTS (SELECT * FROM "
<< buildTableName << " WHERE " << joinConditionAsSql(*joinNode);
sql << ")";
}
break;
Expand All @@ -448,6 +500,17 @@ std::optional<std::string> DuckQueryRunner::toSql(

std::optional<std::string> DuckQueryRunner::toSql(
const std::shared_ptr<const core::NestedLoopJoinNode>& joinNode) {
std::string probeTableName;
std::string buildTableName;
if (const auto probeSubQuery = toSql(joinNode->sources()[0]),
buildSubQuery = toSql(joinNode->sources()[1]);
probeSubQuery && buildSubQuery) {
probeTableName = fmt::format("({})", *probeSubQuery);
buildTableName = fmt::format("({})", *buildSubQuery);
} else {
return std::nullopt;
}

std::stringstream sql;
sql << "SELECT " << folly::join(", ", joinNode->outputType()->names());

Expand All @@ -458,13 +521,16 @@ std::optional<std::string> DuckQueryRunner::toSql(
const std::string joinCondition{"(1 = 1)"};
switch (joinNode->joinType()) {
case core::JoinType::kInner:
sql << " FROM t INNER JOIN u ON " << joinCondition;
sql << " FROM " << probeTableName << " INNER JOIN " << buildTableName
<< " ON " << joinCondition;
break;
case core::JoinType::kLeft:
sql << " FROM t LEFT JOIN u ON " << joinCondition;
sql << " FROM " << probeTableName << " LEFT JOIN " << buildTableName
<< " ON " << joinCondition;
break;
case core::JoinType::kFull:
sql << " FROM t FULL OUTER JOIN u ON " << joinCondition;
sql << " FROM " << probeTableName << " FULL OUTER JOIN " << buildTableName
<< " ON " << joinCondition;
break;
default:
VELOX_UNREACHABLE(
Expand Down
12 changes: 12 additions & 0 deletions velox/exec/fuzzer/DuckQueryRunner.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,18 @@ class DuckQueryRunner : public ReferenceQueryRunner {
/// Assumes that source of AggregationNode or Window Node is 'tmp' table.
std::optional<std::string> toSql(const core::PlanNodePtr& plan) override;

/// Returns the name of the values node table in the form t_<id>.
std::string getTableName(const core::ValuesNodePtr& valuesNode);

// Traverses all nodes in the plan and returns all tables and their names.
std::unordered_map<std::string, std::vector<velox::RowVectorPtr>>
getAllTablesAndNames(const core::PlanNodePtr& plan);

/// Executes SQL query returned by the 'toSql' method based on the plan.
std::multiset<std::vector<velox::variant>> execute(
const std::string& sql,
const core::PlanNodePtr& plan) override;

/// Creates 'tmp' table with 'input' data and runs 'sql' query. Returns
/// results according to 'resultType' schema.
std::multiset<std::vector<velox::variant>> execute(
Expand Down
4 changes: 1 addition & 3 deletions velox/exec/fuzzer/JoinFuzzer.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -680,10 +680,8 @@ std::optional<MaterializedRowMultiset> JoinFuzzer::computeReferenceResults(
}

if (auto sql = referenceQueryRunner_->toSql(plan)) {
return referenceQueryRunner_->execute(
sql.value(), probeInput, buildInput, plan->outputType());
return referenceQueryRunner_->execute(*sql, plan);
}

LOG(INFO) << "Query not supported by the reference DB";
return std::nullopt;
}
Expand Down
Loading

0 comments on commit e222e6d

Please sign in to comment.