implemented compare between

nickgerrets · nickgerrets · commit af0b6f26b39d · 2024-09-04T16:53:09.000+02:00
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
@@ -414,6 +414,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>();
@@ -538,6 +561,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);
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');
diff --git a/test/sql/test_substrait_tpcds.test b/test/sql/test_substrait_tpcds.test
@@ -20,9 +20,9 @@ CALL dsdgen(sf=0.01)
 #statement ok
 #CALL get_substrait('SELECT cd_gender, cd_marital_status, cd_education_status, Count(*) cnt1, cd_purchase_estimate, Count(*) cnt2, cd_credit_rating, Count(*) cnt3, cd_dep_count, Count(*) cnt4, cd_dep_employed_count, Count(*) cnt5, cd_dep_college_count, Count(*) cnt6 FROM customer c, customer_address ca, customer_demographics WHERE c.c_current_addr_sk = ca.ca_address_sk AND ca_county IN ( ''Lycoming County'', ''Sheridan County'', ''Kandiyohi County'', ''Pike County'', ''Greene County'' ) AND cd_demo_sk = c.c_current_cdemo_sk AND EXISTS (SELECT * FROM store_sales, date_dim WHERE c.c_customer_sk = ss_customer_sk AND ss_sold_date_sk = d_date_sk AND d_year = 2002 AND d_moy BETWEEN 4 AND 4 + 3) AND ( EXISTS (SELECT * FROM web_sales, date_dim WHERE c.c_customer_sk = ws_bill_customer_sk AND ws_sold_date_sk = d_date_sk AND d_year = 2002 AND d_moy BETWEEN 4 AND 4 + 3) OR EXISTS (SELECT * FROM catalog_sales, date_dim WHERE c.c_customer_sk = cs_ship_customer_sk AND cs_sold_date_sk = d_date_sk AND d_year = 2002 AND d_moy BETWEEN 4 AND 4 + 3) ) GROUP BY cd_gender, cd_marital_status, cd_education_status, cd_purchase_estimate, cd_credit_rating, cd_dep_count, cd_dep_employed_count, cd_dep_college_count ORDER BY cd_gender, cd_marital_status, cd_education_status, cd_purchase_estimate, cd_credit_rating, cd_dep_count, cd_dep_employed_count, cd_dep_college_count LIMIT 100; ')
 
-#Q 3 (COMPARE_BETWEEN)
-#statement ok
-#CALL get_substrait('WITH year_total AS (SELECT c_customer_id customer_id, c_first_name customer_first_name , c_last_name customer_last_name, c_preferred_cust_flag customer_preferred_cust_flag , c_birth_country customer_birth_country, c_login customer_login, c_email_address customer_email_address, d_year dyear, Sum(ss_ext_list_price - ss_ext_discount_amt) year_total, ''s'' sale_type FROM customer, store_sales, date_dim WHERE c_customer_sk = ss_customer_sk AND ss_sold_date_sk = d_date_sk GROUP BY c_customer_id, c_first_name, c_last_name, c_preferred_cust_flag, c_birth_country, c_login, c_email_address, d_year UNION ALL SELECT c_customer_id customer_id, c_first_name customer_first_name , c_last_name customer_last_name, c_preferred_cust_flag customer_preferred_cust_flag , c_birth_country customer_birth_country, c_login customer_login, c_email_address customer_email_address, d_year dyear, Sum(ws_ext_list_price - ws_ext_discount_amt) year_total, ''w'' sale_type FROM customer, web_sales, date_dim WHERE c_customer_sk = ws_bill_customer_sk AND ws_sold_date_sk = d_date_sk GROUP BY c_customer_id, c_first_name, c_last_name, c_preferred_cust_flag, c_birth_country, c_login, c_email_address, d_year) SELECT t_s_secyear.customer_id, t_s_secyear.customer_first_name, t_s_secyear.customer_last_name, t_s_secyear.customer_birth_country FROM year_total t_s_firstyear, year_total t_s_secyear, year_total t_w_firstyear, year_total t_w_secyear WHERE t_s_secyear.customer_id = t_s_firstyear.customer_id AND t_s_firstyear.customer_id = t_w_secyear.customer_id AND t_s_firstyear.customer_id = t_w_firstyear.customer_id AND t_s_firstyear.sale_type = ''s'' AND t_w_firstyear.sale_type = ''w'' AND t_s_secyear.sale_type = ''s'' AND t_w_secyear.sale_type = ''w'' AND t_s_firstyear.dyear = 2001 AND t_s_secyear.dyear = 2001 + 1 AND t_w_firstyear.dyear = 2001 AND t_w_secyear.dyear = 2001 + 1 AND t_s_firstyear.year_total > 0 AND t_w_firstyear.year_total > 0 AND CASE WHEN t_w_firstyear.year_total > 0 THEN t_w_secyear.year_total / t_w_firstyear.year_total ELSE 0.0 END > CASE WHEN t_s_firstyear.year_total > 0 THEN t_s_secyear.year_total / t_s_firstyear.year_total ELSE 0.0 END ORDER BY t_s_secyear.customer_id, t_s_secyear.customer_first_name, t_s_secyear.customer_last_name, t_s_secyear.customer_birth_country LIMIT 100; ')
+#Q 3
+statement ok
+CALL get_substrait('WITH year_total AS (SELECT c_customer_id customer_id, c_first_name customer_first_name , c_last_name customer_last_name, c_preferred_cust_flag customer_preferred_cust_flag , c_birth_country customer_birth_country, c_login customer_login, c_email_address customer_email_address, d_year dyear, Sum(ss_ext_list_price - ss_ext_discount_amt) year_total, ''s'' sale_type FROM customer, store_sales, date_dim WHERE c_customer_sk = ss_customer_sk AND ss_sold_date_sk = d_date_sk GROUP BY c_customer_id, c_first_name, c_last_name, c_preferred_cust_flag, c_birth_country, c_login, c_email_address, d_year UNION ALL SELECT c_customer_id customer_id, c_first_name customer_first_name , c_last_name customer_last_name, c_preferred_cust_flag customer_preferred_cust_flag , c_birth_country customer_birth_country, c_login customer_login, c_email_address customer_email_address, d_year dyear, Sum(ws_ext_list_price - ws_ext_discount_amt) year_total, ''w'' sale_type FROM customer, web_sales, date_dim WHERE c_customer_sk = ws_bill_customer_sk AND ws_sold_date_sk = d_date_sk GROUP BY c_customer_id, c_first_name, c_last_name, c_preferred_cust_flag, c_birth_country, c_login, c_email_address, d_year) SELECT t_s_secyear.customer_id, t_s_secyear.customer_first_name, t_s_secyear.customer_last_name, t_s_secyear.customer_birth_country FROM year_total t_s_firstyear, year_total t_s_secyear, year_total t_w_firstyear, year_total t_w_secyear WHERE t_s_secyear.customer_id = t_s_firstyear.customer_id AND t_s_firstyear.customer_id = t_w_secyear.customer_id AND t_s_firstyear.customer_id = t_w_firstyear.customer_id AND t_s_firstyear.sale_type = ''s'' AND t_w_firstyear.sale_type = ''w'' AND t_s_secyear.sale_type = ''s'' AND t_w_secyear.sale_type = ''w'' AND t_s_firstyear.dyear = 2001 AND t_s_secyear.dyear = 2001 + 1 AND t_w_firstyear.dyear = 2001 AND t_w_secyear.dyear = 2001 + 1 AND t_s_firstyear.year_total > 0 AND t_w_firstyear.year_total > 0 AND CASE WHEN t_w_firstyear.year_total > 0 THEN t_w_secyear.year_total / t_w_firstyear.year_total ELSE 0.0 END > CASE WHEN t_s_firstyear.year_total > 0 THEN t_s_secyear.year_total / t_s_firstyear.year_total ELSE 0.0 END ORDER BY t_s_secyear.customer_id, t_s_secyear.customer_first_name, t_s_secyear.customer_last_name, t_s_secyear.customer_birth_country LIMIT 100; ')
 
 #Q 4 (WINDOW)
 #statement ok
