Merge pull request #12271 from hawkfish/approx-array

Issue #12252: APPROX_QUANTILE Array Argument
duckdb · May 28, 2024 · cc52145 · cc52145
2 parents ee197d8 + c3927b3
commit cc52145
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 8 deletions.
diff --git a/src/core_functions/aggregate/holistic/approximate_quantile.cpp b/src/core_functions/aggregate/holistic/approximate_quantile.cpp
@@ -181,16 +181,25 @@ unique_ptr<FunctionData> BindApproxQuantile(ClientContext &context, AggregateFun
 		throw BinderException("APPROXIMATE QUANTILE can only take constant quantile parameters");
 	}
 	Value quantile_val = ExpressionExecutor::EvaluateScalar(context, *arguments[1]);
+	if (quantile_val.IsNull()) {
+		throw BinderException("APPROXIMATE QUANTILE parameter list cannot be NULL");
+	}
 
 	vector<float> quantiles;
-	if (quantile_val.type().id() != LogicalTypeId::LIST) {
-		quantiles.push_back(CheckApproxQuantile(quantile_val));
-	} else if (quantile_val.IsNull()) {
-		throw BinderException("APPROXIMATE QUANTILE parameter list cannot be NULL");
-	} else {
+	switch (quantile_val.type().id()) {
+	case LogicalTypeId::LIST:
 		for (const auto &element_val : ListValue::GetChildren(quantile_val)) {
 			quantiles.push_back(CheckApproxQuantile(element_val));
 		}
+		break;
+	case LogicalTypeId::ARRAY:
+		for (const auto &element_val : ArrayValue::GetChildren(quantile_val)) {
+			quantiles.push_back(CheckApproxQuantile(element_val));
+		}
+		break;
+	default:
+		quantiles.push_back(CheckApproxQuantile(quantile_val));
+		break;
 	}
 
 	// remove the quantile argument so we can use the unary aggregate

diff --git a/src/core_functions/aggregate/holistic/quantile.cpp b/src/core_functions/aggregate/holistic/quantile.cpp
@@ -1509,12 +1509,20 @@ unique_ptr<FunctionData> BindQuantile(ClientContext &context, AggregateFunction
 		throw BinderException("QUANTILE argument must not be NULL");
 	}
 	vector<Value> quantiles;
-	if (quantile_val.type().id() != LogicalTypeId::LIST) {
-		quantiles.push_back(CheckQuantile(quantile_val));
-	} else {
+	switch (quantile_val.type().id()) {
+	case LogicalTypeId::LIST:
 		for (const auto &element_val : ListValue::GetChildren(quantile_val)) {
 			quantiles.push_back(CheckQuantile(element_val));
 		}
+		break;
+	case LogicalTypeId::ARRAY:
+		for (const auto &element_val : ArrayValue::GetChildren(quantile_val)) {
+			quantiles.push_back(CheckQuantile(element_val));
+		}
+		break;
+	default:
+		quantiles.push_back(CheckQuantile(quantile_val));
+		break;
 	}
 
 	Function::EraseArgument(function, arguments, arguments.size() - 1);

diff --git a/test/sql/aggregate/aggregates/test_approx_quantile.test b/test/sql/aggregate/aggregates/test_approx_quantile.test
@@ -141,6 +141,19 @@ FROM (
 ----
 [true, true, true]
 
+# Array lists
+query I
+SELECT approx_quantile(col, [0.5, 0.4, 0.1]) AS percentile 
+FROM VALUES (0), (1), (2), (10) AS tab(col);
+----
+[2, 1, 0]
+
+query I
+SELECT approx_quantile(col, ARRAY_VALUE(0.5, 0.4, 0.1)) AS percentile 
+FROM VALUES (0), (1), (2), (10) AS tab(col);
+----
+[2, 1, 0]
+
 # Errors
 statement error
 SELECT approx_quantile(r, -0.1) FROM quantile

diff --git a/test/sql/aggregate/aggregates/test_quantile_disc_list.test b/test/sql/aggregate/aggregates/test_quantile_disc_list.test
@@ -165,6 +165,13 @@ SELECT quantile_disc(42::UTINYINT, 0.5);
 ----
 42
 
+# Array arguments
+query I
+SELECT quantile_disc(col, ARRAY_VALUE(0.5, 0.4, 0.1)) AS percentile 
+FROM VALUES (0), (1), (2), (10) AS tab(col);
+----
+[1, 1, 0]
+
 # Invalid use
 statement error
 SELECT quantile_disc(r, [-0.1, 0.5, 0.9]) FROM quantiles