refactor: query config options (#6781)

* feat: refactor columnar and vector conversion

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* feat: initialize config options from query context

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* fix: failure tests

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* chore: revert ColumnarValue::try_from_vector

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

---------

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>
This commit is contained in:
dennis zhuang
2025-09-01 15:00:26 +08:00
committed by GitHub
parent d57c0db9e6
commit 1234911ed3
16 changed files with 103 additions and 69 deletions

View File

@@ -19,6 +19,7 @@ use common_error::define_into_tonic_status;
use common_error::ext::{BoxedError, ErrorExt};
use common_error::status_code::StatusCode;
use common_macro::stack_trace_debug;
use common_query::error::Error as QueryResult;
use datafusion::parquet;
use datafusion_common::DataFusionError;
use datatypes::arrow::error::ArrowError;
@@ -36,6 +37,14 @@ pub enum Error {
location: Location,
},
#[snafu(display("Failed to cast result: `{}`", source))]
Cast {
#[snafu(source)]
source: QueryResult,
#[snafu(implicit)]
location: Location,
},
#[snafu(display("View already exists: `{name}`"))]
ViewAlreadyExists {
name: String,
@@ -870,6 +879,7 @@ pub type Result<T> = std::result::Result<T, Error>;
impl ErrorExt for Error {
fn status_code(&self) -> StatusCode {
match self {
Error::Cast { source, .. } => source.status_code(),
Error::InvalidSql { .. }
| Error::InvalidConfigValue { .. }
| Error::InvalidInsertRequest { .. }

View File

@@ -32,7 +32,7 @@ use snafu::{ensure, OptionExt, ResultExt};
use sql::ast::{Expr, FunctionArg, FunctionArgExpr, FunctionArguments, Value as SqlValue};
use sql::statements::admin::Admin;
use crate::error::{self, ExecuteAdminFunctionSnafu, IntoVectorsSnafu, Result};
use crate::error::{self, CastSnafu, ExecuteAdminFunctionSnafu, Result};
use crate::statement::StatementExecutor;
const DUMMY_COLUMN: &str = "<dummy>";
@@ -118,7 +118,7 @@ impl StatementExecutor {
.collect(),
return_field: Arc::new(arrow::datatypes::Field::new("result", ret_type, true)),
number_rows: if args.is_empty() { 1 } else { args[0].len() },
config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
config_options: Arc::new(query_ctx.create_config_options()),
};
// Execute the async UDF
@@ -134,22 +134,11 @@ impl StatementExecutor {
})?;
// Convert result back to VectorRef
let result = match result_columnar {
datafusion_expr::ColumnarValue::Array(array) => {
datatypes::vectors::Helper::try_into_vector(array).context(IntoVectorsSnafu)?
}
datafusion_expr::ColumnarValue::Scalar(scalar) => {
let array =
scalar
.to_array_of_size(1)
.with_context(|_| ExecuteAdminFunctionSnafu {
msg: format!("Failed to convert scalar to array for {}", fn_name),
})?;
datatypes::vectors::Helper::try_into_vector(array).context(IntoVectorsSnafu)?
}
};
let result_columnar: common_query::prelude::ColumnarValue =
(&result_columnar).try_into().context(CastSnafu)?;
let result_vector: VectorRef = result_columnar.try_into_vector(1).context(CastSnafu)?;
let result_vector: VectorRef = result;
let column_schemas = vec![ColumnSchema::new(
// Use statement as the result column name
stmt.to_string(),