Alright delim join get working with refrels

pdet · Aug 2, 2024 · d6af7b6 · d6af7b6
1 parent bde22f4
commit d6af7b6
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 13 deletions.
diff --git a/src/from_substrait.cpp b/src/from_substrait.cpp
@@ -439,7 +439,6 @@ shared_ptr<Relation> SubstraitToDuckDB::TransformDelimJoinOp(const substrait::Re
 		duplicate_eliminated_columns.emplace_back(
 		    make_uniq<PositionalReferenceExpression>(col.direct_reference().struct_field().field() + 1));
 	}
-	duplicate_eliminated_columns_ptr = &duplicate_eliminated_columns;
 
 	JoinType djointype;
 	switch (sjoin.type()) {
@@ -484,18 +483,14 @@ shared_ptr<Relation> SubstraitToDuckDB::TransformDelimJoinOp(const substrait::Re
 }
 
 shared_ptr<Relation> SubstraitToDuckDB::TransformDelimGetOp(const substrait::Rel &sop) {
-	auto &delimiter_get = sop.delim_get();
-	auto subtree = TransformReferenceOp(delimiter_get.input());
+	auto &delim_get = sop.delim_get();
+	auto subtree = TransformReferenceOp(delim_get.input());
 
-	// vector<unique_ptr<ParsedExpression>> expressions;
-	// for (auto &sexpr : sop.project().expressions()) {
-	// 	expressions.push_back(TransformExpr(sexpr));
-	// }
 	auto &client_context = con.context;
 	vector<LogicalType> chunk_types;
 	auto &input_columns = subtree->Columns();
-	for (auto &col_ref : *duplicate_eliminated_columns_ptr) {
-		chunk_types.emplace_back(input_columns[col_ref->Cast<PositionalReferenceExpression>().index - 1].Type());
+	for (auto &col_ref : delim_get.column_ids()) {
+		chunk_types.emplace_back(input_columns[col_ref.direct_reference().struct_field().field()].Type());
 	}
 	return make_shared_ptr<DelimGetRelation>(client_context, chunk_types);
 }

diff --git a/src/include/from_substrait.hpp b/src/include/from_substrait.hpp
@@ -60,8 +60,5 @@ class SubstraitToDuckDB {
 	//! names
 	static const unordered_map<std::string, std::string> function_names_remap;
 	static const case_insensitive_set_t valid_extract_subfields;
-
-	//! FIXME: this is pretty brittle
-	vector<unique_ptr<ParsedExpression>> *duplicate_eliminated_columns_ptr = nullptr;
 };
 } // namespace duckdb
diff --git a/src/include/to_substrait.hpp b/src/include/to_substrait.hpp
@@ -170,5 +170,7 @@ class DuckDBToSubstrait {
 	//! This really only matters for delim joins/gets, since these are
 	//! the only splits we currently support.
 	int32_t cur_subtree_relation = 1;
+	//! The pointer to a delim join
+	LogicalComparisonJoin *duplicate_eliminated_parent_ptr = nullptr;
 };
 } // namespace duckdb
diff --git a/src/to_substrait.cpp b/src/to_substrait.cpp
@@ -941,10 +941,12 @@ substrait::Rel *DuckDBToSubstrait::TransformComparisonJoin(LogicalOperator &dop)
 }
 
 substrait::Rel *DuckDBToSubstrait::TransformDelimiterJoin(LogicalOperator &dop) {
+	auto &djoin = dop.Cast<LogicalComparisonJoin>();
+	duplicate_eliminated_parent_ptr = &djoin;
 	auto res = new substrait::Rel();
 
 	auto sjoin = res->mutable_delim_join();
-	auto &djoin = (LogicalComparisonJoin &)dop;
+
 	auto lhs_child = TransformOp(*dop.children[0]);
 	auto rhs_child = TransformOp(*dop.children[1]);
 	if (djoin.delim_flipped) {
@@ -1388,8 +1390,15 @@ substrait::Rel *DuckDBToSubstrait::TransformIntersect(LogicalOperator &dop) {
 substrait::Rel *DuckDBToSubstrait::TransformDelimGet() {
 	auto rel = new substrait::Rel();
 	auto delim_get = rel->mutable_delim_get();
+	D_ASSERT(duplicate_eliminated_parent_ptr);
+	for (auto &dup_col : duplicate_eliminated_parent_ptr->duplicate_eliminated_columns) {
+		auto &dref = dup_col->Cast<BoundReferenceExpression>();
+		auto s_dup_col = delim_get->add_column_ids();
+		s_dup_col->mutable_direct_reference()->mutable_struct_field()->set_field(static_cast<int32_t>(dref.index));
+	}
 	auto ref_input = delim_get->mutable_input();
 	ref_input->set_subtree_ordinal(cur_subtree_relation);
+	duplicate_eliminated_parent_ptr = nullptr;
 	return rel;
 }