Merge pull request #104 from pdet/nested_expressions

Nested expressions

Author: pdet

src/from_substrait.cpp

@@ -325,6 +325,37 @@ unique_ptr<ParsedExpression> SubstraitToDuckDB::TransformInExpr(const substrait:
 	return make_uniq<OperatorExpression>(ExpressionType::COMPARE_IN, std::move(values));
 }
 
+unique_ptr<ParsedExpression> SubstraitToDuckDB::TransformNested(const substrait::Expression &sexpr) {
+	auto &nested_expression = sexpr.nested();
+	if (nested_expression.has_struct_()) {
+		auto &struct_expression = nested_expression.struct_();
+		 vector<unique_ptr<ParsedExpression>> children;
+		for (auto& child: struct_expression.fields()) {
+			children.emplace_back(TransformExpr(child));
+		}
+		return make_uniq<FunctionExpression>("row", std::move(children));
+	} else if (nested_expression.has_list()) {
+		auto &list_expression = nested_expression.list();
+		 vector<unique_ptr<ParsedExpression>> children;
+		for (auto& child: list_expression.values()) {
+			children.emplace_back(TransformExpr(child));
+		}
+		return make_uniq<FunctionExpression>("list_value", std::move(children));
+
+	} else if (nested_expression.has_map()) {
+		auto &map_expression = nested_expression.map();
+		 vector<unique_ptr<ParsedExpression>> children;
+		auto key_value = map_expression.key_values();
+		children.emplace_back(TransformExpr(key_value[0].key()));
+		children.emplace_back(TransformExpr(key_value[0].value()));
+		return make_uniq<FunctionExpression>("map", std::move(children));
+
+	} else{
+		throw NotImplementedException("Substrait nested expression is not yet implemented.");
+	}
+
+}
+
 unique_ptr<ParsedExpression> SubstraitToDuckDB::TransformExpr(const substrait::Expression &sexpr) {
 	switch (sexpr.rex_type_case()) {
 	case substrait::Expression::RexTypeCase::kLiteral:
@@ -339,6 +370,8 @@ unique_ptr<ParsedExpression> SubstraitToDuckDB::TransformExpr(const substrait::E
 		return TransformCastExpr(sexpr);
 	case substrait::Expression::RexTypeCase::kSingularOrList:
 		return TransformInExpr(sexpr);
+	case substrait::Expression::RexTypeCase::kNested:
+		return TransformNested(sexpr);
 	case substrait::Expression::RexTypeCase::kSubquery:
 	default:
 		throw InternalException("Unsupported expression type " + to_string(sexpr.rex_type_case()));

src/include/from_substrait.hpp

@@ -37,6 +37,7 @@ class SubstraitToDuckDB {
 	unique_ptr<ParsedExpression> TransformIfThenExpr(const substrait::Expression &sexpr);
 	unique_ptr<ParsedExpression> TransformCastExpr(const substrait::Expression &sexpr);
 	unique_ptr<ParsedExpression> TransformInExpr(const substrait::Expression &sexpr);
+	unique_ptr<ParsedExpression> TransformNested(const substrait::Expression &sexpr);
 
 	static void VerifyCorrectExtractSubfield(const string &subfield);
 	static string RemapFunctionName(const string &function_name);
@@ -57,5 +58,6 @@ class SubstraitToDuckDB {
 	//! names
 	static const unordered_map<std::string, std::string> function_names_remap;
 	static const case_insensitive_set_t valid_extract_subfields;
+	vector<ParsedExpression*> struct_expressions;
 };
 } // namespace duckdb

src/to_substrait.cpp

@@ -314,17 +314,48 @@ bool DuckDBToSubstrait::IsExtractFunction(const string &function_name) {
 void DuckDBToSubstrait::TransformFunctionExpression(Expression &dexpr, substrait::Expression &sexpr,
                                                     uint64_t col_offset) {
 	auto &dfun = dexpr.Cast<BoundFunctionExpression>();
-	auto sfun = sexpr.mutable_scalar_function();
+
 
 	auto function_name = dfun.function.name;
+
+	if (function_name == "row") {
+		auto nested_expression = sexpr.mutable_nested();
+		auto struct_expression = nested_expression->mutable_struct_();
+		for (auto& child: dfun.children) {
+			auto child_expression = struct_expression->add_fields();
+			TransformExpr(*child, *child_expression);
+		}
+		return;
+	}
+	if (function_name == "list_value" || function_name == "list_pack") {
+		auto nested_expression = sexpr.mutable_nested();
+		auto list_expression = nested_expression->mutable_list();
+		for (auto& child: dfun.children) {
+			auto child_value = list_expression->add_values();
+			TransformExpr(*child, *child_value);
+		}
+		return;
+	}
+	if (function_name == "map") {
+		auto nested_expression = sexpr.mutable_nested();
+		auto map_expression = nested_expression->mutable_map();
+		D_ASSERT(dfun.children.size() == 2);
+		auto child_value = map_expression->add_key_values();
+		auto key = child_value->mutable_key();
+		auto value = child_value->mutable_value();
+		TransformExpr(*dfun.children[0], *key);
+		TransformExpr(*dfun.children[1], *value);
+		return;
+	}
+	auto sfun = sexpr.mutable_scalar_function();
 	if (IsExtractFunction(function_name)) {
 		// Change the name to 'extract', and add an Enum argument containing the subfield
 		auto subfield = function_name;
 		function_name = "extract";
 		auto enum_arg = sfun->add_arguments();
 		*enum_arg->mutable_enum_() = subfield;
 	}
-	vector<::substrait::Type> args_types;
+	vector<substrait::Type> args_types;
 	for (auto &darg : dfun.children) {
 		auto sarg = sfun->add_arguments();
 		TransformExpr(*darg, *sarg->mutable_value(), col_offset);

test/sql/test_nested_expressions.test

@@ -0,0 +1,48 @@
+# name: test/sql/test_nested_expressions.test
+# description: Test nested expressions
+# group: [sql]
+
+require substrait
+
+statement ok
+PRAGMA enable_verification
+
+statement ok
+create table t as select range as a from range(10)
+
+# Test struct creation
+statement ok
+CALL get_substrait('SELECT row(a,a,10) from t;')
+
+# Test list creation
+statement ok
+CALL get_substrait('SELECT [a,a,10] from t;')
+
+# Test map creation
+statement ok
+CALL get_substrait('SELECT MAP {a: a} from t;')
+
+# Test nested-> nested
+statement ok
+CALL get_substrait('SELECT MAP {[a,a,10]: [a,a,10]} from t;')
+
+statement ok
+CALL get_substrait('SELECT MAP {[row(a,a,10),row(a,a,10),row(a,a,10)]: [row(a,a,10),row(a,a,10),row(a,a,10)]} from t;')
+
+statement ok
+CALL get_substrait('SELECT row([a,a,10],[a,a,10],[a,a,10]) from t;')
+
+statement ok
+CALL get_substrait('SELECT row([MAP {a: a},MAP {a: a},MAP {a: a}],[MAP {a: a},MAP {a: a},MAP {a: a}],[MAP {a: a},MAP {a: a},MAP {a: a}]) from t;')
+
+statement ok
+CALL get_substrait('SELECT MAP {[a,a,10]: [a,a,10]} from t;')
+
+statement ok
+CALL get_substrait('SELECT MAP {[row(a,a,10),row(a,a,10),row(a,a,10)]: [row(a,a,10),row(a,a,10),row(a,a,10)]} from t;')
+
+statement ok
+CALL get_substrait('SELECT row(row(row(a,a,10),row(a,a,10),row(a,a,10)),row(row(a,a,10),row(a,a,10),row(a,a,10)),row(row(a,a,10),row(a,a,10),row(a,a,10))) from t;')
+
+statement ok
+CALL get_substrait('SELECT [[[a,a,10], [a,a,10], [a,a,10]], [[a,a,10], [a,a,10], [a,a,10]], [[a,a,10], [a,a,10], [a,a,10]]] from t;')