Merge pull request #106 from nickgerrets/tpcds

TPC-DS queries and BETWEEN comparison
substrait-io · Sep 18, 2024 · 78d265c · 78d265c
2 parents d8ca65b + 5ef934a
commit 78d265c
Show file tree

Hide file tree

Showing 7 changed files with 579 additions and 9 deletions.
diff --git a/duckdb-r b/duckdb-r
diff --git a/src/from_substrait.cpp b/src/from_substrait.cpp
@@ -254,7 +254,6 @@ unique_ptr<ParsedExpression> SubstraitToDuckDB::TransformScalarFunctionExpr(cons
 		return make_uniq<ComparisonExpression>(ExpressionType::COMPARE_NOT_DISTINCT_FROM, std::move(children[0]),
 		                                       std::move(children[1]));
 	} else if (function_name == "between") {
-		// FIXME: ADD between to substrait extension
 		D_ASSERT(children.size() == 3);
 		return make_uniq<BetweenExpression>(std::move(children[0]), std::move(children[1]), std::move(children[2]));
 	} else if (function_name == "extract") {

diff --git a/src/include/to_substrait.hpp b/src/include/to_substrait.hpp
@@ -88,6 +88,7 @@ class DuckDBToSubstrait {
 	void TransformFunctionExpression(Expression &dexpr, substrait::Expression &sexpr, uint64_t col_offset);
 	static void TransformConstantExpression(Expression &dexpr, substrait::Expression &sexpr);
 	void TransformComparisonExpression(Expression &dexpr, substrait::Expression &sexpr);
+	void TransformBetweenExpression(Expression &dexpr, substrait::Expression &sexpr);
 	void TransformConjunctionExpression(Expression &dexpr, substrait::Expression &sexpr, uint64_t col_offset);
 	void TransformNotNullExpression(Expression &dexpr, substrait::Expression &sexpr, uint64_t col_offset);
 	void TransformIsNullExpression(Expression &dexpr, substrait::Expression &sexpr, uint64_t col_offset);

diff --git a/src/to_substrait.cpp b/src/to_substrait.cpp
@@ -413,6 +413,29 @@ void DuckDBToSubstrait::TransformComparisonExpression(Expression &dexpr, substra
 	*scalar_fun->mutable_output_type() = DuckToSubstraitType(dcomp.return_type);
 }
 
+void DuckDBToSubstrait::TransformBetweenExpression(Expression &dexpr, substrait::Expression &sexpr) {
+	auto &dcomp = dexpr.Cast<BoundBetweenExpression>();
+
+	if (dexpr.type != ExpressionType::COMPARE_BETWEEN) {
+		throw InternalException("Not a between comparison expression");
+	}
+
+	auto scalar_fun = sexpr.mutable_scalar_function();
+	vector<::substrait::Type> args_types;
+	args_types.emplace_back(DuckToSubstraitType(dcomp.input->return_type));
+	args_types.emplace_back(DuckToSubstraitType(dcomp.lower->return_type));
+	args_types.emplace_back(DuckToSubstraitType(dcomp.upper->return_type));
+	scalar_fun->set_function_reference(RegisterFunction("between", args_types));
+
+	auto sarg = scalar_fun->add_arguments();
+	TransformExpr(*dcomp.input, *sarg->mutable_value(), 0);
+	sarg = scalar_fun->add_arguments();
+	TransformExpr(*dcomp.lower, *sarg->mutable_value(), 0);
+	sarg = scalar_fun->add_arguments();
+	TransformExpr(*dcomp.upper, *sarg->mutable_value(), 0);
+	*scalar_fun->mutable_output_type() = DuckToSubstraitType(dcomp.return_type);
+}
+
 void DuckDBToSubstrait::TransformConjunctionExpression(Expression &dexpr, substrait::Expression &sexpr,
                                                        uint64_t col_offset) {
 	auto &dconj = dexpr.Cast<BoundConjunctionExpression>();
@@ -537,6 +560,9 @@ void DuckDBToSubstrait::TransformExpr(Expression &dexpr, substrait::Expression &
 	case ExpressionType::COMPARE_NOT_DISTINCT_FROM:
 		TransformComparisonExpression(dexpr, sexpr);
 		break;
+	case ExpressionType::COMPARE_BETWEEN:
+		TransformBetweenExpression(dexpr, sexpr);
+		break;
 	case ExpressionType::CONJUNCTION_AND:
 	case ExpressionType::CONJUNCTION_OR:
 		TransformConjunctionExpression(dexpr, sexpr, col_offset);
@@ -557,7 +583,7 @@ void DuckDBToSubstrait::TransformExpr(Expression &dexpr, substrait::Expression &
 		TransformNotExpression(dexpr, sexpr, col_offset);
 		break;
 	default:
-		throw InternalException(ExpressionTypeToString(dexpr.type));
+		throw NotImplementedException(ExpressionTypeToString(dexpr.type));
 	}
 }
 
@@ -742,7 +768,7 @@ substrait::Expression *DuckDBToSubstrait::TransformJoinCond(const JoinCondition
 		join_comparision = "lt";
 		break;
 	default:
-		throw InternalException("Unsupported join comparison: " + ExpressionTypeToOperator(dcond.comparison));
+		throw NotImplementedException("Unsupported join comparison: " + ExpressionTypeToOperator(dcond.comparison));
 	}
 	vector<::substrait::Type> args_types;
 	auto scalar_fun = expr->mutable_scalar_function();
@@ -946,7 +972,7 @@ substrait::Rel *DuckDBToSubstrait::TransformComparisonJoin(LogicalOperator &dop)
 		sjoin->set_type(substrait::JoinRel::JoinType::JoinRel_JoinType_JOIN_TYPE_OUTER);
 		break;
 	default:
-		throw InternalException("Unsupported join type " + JoinTypeToString(djoin.join_type));
+		throw NotImplementedException("Unsupported join type " + JoinTypeToString(djoin.join_type));
 	}
 	// somewhat odd semantics on our side
 	if (djoin.left_projection_map.empty()) {
@@ -984,15 +1010,15 @@ substrait::Rel *DuckDBToSubstrait::TransformAggregateGroup(LogicalOperator &dop)
 	for (auto &dgrp : daggr.groups) {
 		if (dgrp->type != ExpressionType::BOUND_REF) {
 			// TODO push projection or push substrait to allow expressions here
-			throw InternalException("No expressions in groupings yet");
+			throw NotImplementedException("No expressions in groupings yet");
 		}
 		TransformExpr(*dgrp, *sgrp->add_grouping_expressions());
 	}
 	for (auto &dmeas : daggr.expressions) {
 		auto smeas = saggr->add_measures()->mutable_measure();
 		if (dmeas->type != ExpressionType::BOUND_AGGREGATE) {
 			// TODO push projection or push substrait, too
-			throw InternalException("No non-aggregate expressions in measures yet");
+			throw NotImplementedException("No non-aggregate expressions in measures yet");
 		}
 		auto &daexpr = dmeas->Cast<BoundAggregateExpression>();
 
@@ -1422,7 +1448,7 @@ substrait::Rel *DuckDBToSubstrait::TransformOp(LogicalOperator &dop) {
 	case LogicalOperatorType::LOGICAL_DUMMY_SCAN:
 		return TransformDummyScan();
 	default:
-		throw InternalException(LogicalOperatorToString(dop.type));
+		throw NotImplementedException(LogicalOperatorToString(dop.type));
 	}
 }
 

diff --git a/substrait b/substrait
diff --git a/test/sql/test_between.test b/test/sql/test_between.test
@@ -0,0 +1,14 @@
+# name: test/sql/test_between.test
+# description: Test BETWEEN comparison
+# group: [sql]
+
+require substrait
+
+statement ok
+PRAGMA enable_verification
+
+statement ok
+create table t as select * from range(100) as t(x)
+
+statement ok
+CALL get_substrait('select * from t where x BETWEEN 4 AND 6;', enable_optimizer = false );
+4 −4		.github/workflows/pr.yml
+1 −1		.github/workflows/pr_breaking.yml
+1 −1		.github/workflows/release.yml
+56 −0		CHANGELOG.md
+41 −0		extensions/functions_aggregate_decimal_output.yaml
+73 −1		extensions/functions_arithmetic_decimal.yaml
+266 −35		extensions/functions_datetime.yaml
+59 −0		extensions/functions_logarithmic.yaml
+65 −8		proto/substrait/algebra.proto
+14 −1		proto/substrait/parameterized_types.proto
+15 −0		proto/substrait/type.proto
+14 −1		proto/substrait/type_expressions.proto
+1 −0		site/docs/extensions/index.md
+9 −3		site/docs/relations/basics.md
+11 −6		site/docs/relations/logical_relations.md
+4 −2		site/docs/serialization/_config
+25 −0		site/docs/serialization/basics.md
+1 −0		site/docs/types/_config
+95 −0		site/docs/types/named_structs.md
+4 −3		site/docs/types/type_classes.md
+1 −1		site/requirements.txt