21
21
22
22
#include " duckdb/parser/expression/comparison_expression.hpp"
23
23
24
- #include " substrait/plan.pb.h"
25
- #include " google/protobuf/util/json_util.h"
26
24
#include " duckdb/main/client_data.hpp"
25
+ #include " google/protobuf/util/json_util.h"
26
+ #include " substrait/plan.pb.h"
27
+
28
+ #include " duckdb/main/relation/table_relation.hpp"
27
29
28
30
namespace duckdb {
29
31
const std::unordered_map<std::string, std::string> SubstraitToDuckDB::function_names_remap = {
@@ -661,21 +663,22 @@ int32_t SkipColumnNames(const LogicalType &type) {
661
663
return columns_to_skip;
662
664
}
663
665
664
- Relation *GetProjectionRelation (Relation &relation, string &error) {
666
+ Relation *GetProjectionOrTableRelation (Relation &relation, string &error) {
665
667
error += RelationTypeToString (relation.type );
666
668
switch (relation.type ) {
669
+ case RelationType::TABLE_RELATION:
667
670
case RelationType::PROJECTION_RELATION:
668
671
error += " -> " ;
669
672
return &relation;
670
673
case RelationType::LIMIT_RELATION:
671
674
error += " -> " ;
672
- return GetProjectionRelation (*relation.Cast <LimitRelation>().child , error);
675
+ return GetProjectionOrTableRelation (*relation.Cast <LimitRelation>().child , error);
673
676
case RelationType::ORDER_RELATION:
674
677
error += " -> " ;
675
- return GetProjectionRelation (*relation.Cast <OrderRelation>().child , error);
678
+ return GetProjectionOrTableRelation (*relation.Cast <OrderRelation>().child , error);
676
679
case RelationType::SET_OPERATION_RELATION:
677
680
error += " -> " ;
678
- return GetProjectionRelation (*relation.Cast <SetOpRelation>().right , error);
681
+ return GetProjectionOrTableRelation (*relation.Cast <SetOpRelation>().right , error);
679
682
default :
680
683
throw NotImplementedException (
681
684
" Relation %s is not yet implemented as a possible root chain type of from_substrait function" , error);
@@ -684,15 +687,20 @@ Relation *GetProjectionRelation(Relation &relation, string &error) {
684
687
685
688
shared_ptr<Relation> SubstraitToDuckDB::TransformRootOp (const substrait::RelRoot &sop) {
686
689
vector<string> aliases;
687
- auto column_names = sop.names ();
690
+ const auto & column_names = sop.names ();
688
691
vector<unique_ptr<ParsedExpression>> expressions;
689
692
int id = 1 ;
690
693
auto child = TransformOp (sop.input ());
691
694
string error;
692
- auto first_projection = GetProjectionRelation (*child, error);
693
- auto &columns = first_projection->Cast <ProjectionRelation>().columns ;
695
+ auto first_projection_or_table = GetProjectionOrTableRelation (*child, error);
696
+ vector<ColumnDefinition> *column_definitions;
697
+ if (first_projection_or_table->type == RelationType::PROJECTION_RELATION) {
698
+ column_definitions = &first_projection_or_table->Cast <ProjectionRelation>().columns ;
699
+ } else {
700
+ column_definitions = &first_projection_or_table->Cast <TableRelation>().description ->columns ;
701
+ }
694
702
int32_t i = 0 ;
695
- for (auto &column : columns ) {
703
+ for (auto &column : *column_definitions ) {
696
704
aliases.push_back (column_names[i++]);
697
705
auto column_type = column.GetType ();
698
706
i += SkipColumnNames (column.GetType ());
0 commit comments