Skip to content

Commit fc2ba56

Browse files
Xiong-DaXiong-Da
committed
support get udf in calcite sql
1 parent bf4cb9c commit fc2ba56

File tree

16 files changed

+499
-88
lines changed

16 files changed

+499
-88
lines changed

‎docs/docs/program_model.md‎

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -687,12 +687,35 @@ public class MyTableFunction {
687687

688688
SQLRec 会根据 `eval` 方法的参数类型自动注入相应的值：
689689

690-
| 参数类型 | 注入来源 | SQL 语法 |
691-
|----------|----------|----------|
692-
| `CacheTable` | 传入的缓存表 | 标识符（如 `table_name`） |
693-
| `String` | 字符串字面量或变量 | `'value'` 或 `GET('var')` |
694-
| `ExecuteContext` | 执行上下文 | 自动注入，无需在 SQL 中指定 |
695-
| `ConfigContext` | 配置上下文 | 自动注入，无需在 SQL 中指定 |
690+
| 参数类型 | 注入来源 | SQL 语法 | 适用场景 |
691+
|----------|----------|----------|----------|
692+
| `CacheTable` | 传入的缓存表 | 标识符（如 `table_name`） | 表函数 |
693+
| `String` | 字符串字面量或变量 | `'value'` 或 `GET('var')` | 表函数、标量函数 |
694+
| `ExecuteContext` | 执行上下文 | 自动注入，无需在 SQL 中指定 | 表函数 |
695+
| `ConfigContext` | 配置上下文 | 自动注入，无需在 SQL 中指定 | 表函数 |
696+
| `SqlRecDataContext` | SQLRec 数据上下文 | 自动注入，无需在 SQL 中指定 | 标量函数 |
697+
698+
`SqlRecDataContext` 是专门为标量 UDF 设计的接口，继承自 Calcite 的 `DataContext`。它提供了访问执行上下文变量的能力：
699+
700+
```java
701+
public interface SqlRecDataContext extends DataContext {
702+
String getVariable(String key);
703+
}
704+
```
705+
706+
在标量 UDF 中，可以通过 `SqlRecDataContext` 获取变量值：
707+
708+
```java
709+
public class GetFunction {
710+
public static String eval(DataContext context, String key) {
711+
if (!(context instanceof SqlRecDataContext)) {
712+
throw new IllegalArgumentException("context must be SqlRecDataContext");
713+
}
714+
SqlRecDataContext sqlRecDataContext = (SqlRecDataContext) context;
715+
return sqlRecDataContext.getVariable(key);
716+
}
717+
}
718+
```
696719

697720
参数注入示例
698721

‎docs/docs/udf.md‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -286,6 +286,60 @@ LIMIT 300;
286286
- 如果向量已归一化，内积等于余弦相似度
287287
- 常用于向量检索和相似度计算
288288

289+
---
290+
291+
### get
292+
293+
变量获取函数，从执行上下文中获取变量的值。常用于在SQL中引用通过 `SET` 语句设置的变量。
294+
295+
**函数签名**：
296+
297+
```java
298+
public static String eval(DataContext context, String key)
299+
```
300+
301+
**参数说明**：
302+
303+
| 参数 | 类型 | 说明 |
304+
|------|------|------|
305+
| `key` | String | 变量名 |
306+
307+
**返回值**：返回变量的值（`String`），如果变量不存在则返回 `NULL`。
308+
309+
::: warning 注意
310+
由于 `get` 是 SQL 关键字，使用时需要用反引号包裹函数名，写作 `` `get` ``。
311+
:::
312+
313+
**使用示例**：
314+
315+
```sql
316+
-- 设置变量
317+
SET 'user_id' = '12345';
318+
319+
-- 获取变量值
320+
SELECT `get`('user_id') AS user_id;
321+
322+
-- 在表达式中使用
323+
SELECT `get`('user_id') || '_suffix' AS user_id_with_suffix;
324+
325+
-- 类型转换
326+
SELECT CAST(`get`('limit_count') AS INT) AS limit_count;
327+
328+
-- 从表中获取变量名并使用
329+
CACHE TABLE var_names AS SELECT 'user_id' AS var_name;
330+
SELECT `get`(var_name) AS var_value FROM var_names;
331+
```
332+
333+
**工作原理**：
334+
1. 函数接收一个变量名作为参数
335+
2. 从执行上下文（`ExecuteContext`）中查找对应的变量值
336+
3. 返回变量值，如果变量不存在则返回 `NULL`
337+
338+
**典型应用场景**：
339+
- 参数化SQL查询
340+
- 动态配置传递
341+
- 跨语句共享变量
342+
289343
## 自定义 UDF
290344

291345
可以参考 [编程模型](program_model.md#udf) 文档了解如何开发自定义 UDF。

‎docs/en/docs/program_model.md‎

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -687,12 +687,35 @@ public class MyTableFunction {
687687

688688
SQLRec automatically injects corresponding values based on the `eval` method parameter types:
689689

690-
| Parameter Type | Injection Source | SQL Syntax |
691-
|----------------|------------------|------------|
692-
| `CacheTable` | Passed cache table | Identifier (like `table_name`) |
693-
| `String` | String literal or variable | `'value'` or `GET('var')` |
694-
| `ExecuteContext` | Execution context | Auto-injected, no need to specify in SQL |
695-
| `ConfigContext` | Configuration context | Auto-injected, no need to specify in SQL |
690+
| Parameter Type | Injection Source | SQL Syntax | Use Case |
691+
|----------------|------------------|------------|----------|
692+
| `CacheTable` | Passed cache table | Identifier (like `table_name`) | Table function |
693+
| `String` | String literal or variable | `'value'` or `GET('var')` | Table function, Scalar function |
694+
| `ExecuteContext` | Execution context | Auto-injected, no need to specify in SQL | Table function |
695+
| `ConfigContext` | Configuration context | Auto-injected, no need to specify in SQL | Table function |
696+
| `SqlRecDataContext` | SQLRec data context | Auto-injected, no need to specify in SQL | Scalar function |
697+
698+
`SqlRecDataContext` is an interface specifically designed for scalar UDFs, inheriting from Calcite's `DataContext`. It provides the ability to access execution context variables:
699+
700+
```java
701+
public interface SqlRecDataContext extends DataContext {
702+
String getVariable(String key);
703+
}
704+
```
705+
706+
In scalar UDFs, you can retrieve variable values through `SqlRecDataContext`:
707+
708+
```java
709+
public class GetFunction {
710+
public static String eval(DataContext context, String key) {
711+
if (!(context instanceof SqlRecDataContext)) {
712+
throw new IllegalArgumentException("context must be SqlRecDataContext");
713+
}
714+
SqlRecDataContext sqlRecDataContext = (SqlRecDataContext) context;
715+
return sqlRecDataContext.getVariable(key);
716+
}
717+
}
718+
```
696719

697720
Parameter injection example
698721

‎docs/en/docs/udf.md‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -286,6 +286,60 @@ LIMIT 300;
286286
- If vectors are already normalized, inner product equals cosine similarity
287287
- Commonly used for vector retrieval and similarity calculation
288288

289+
---
290+
291+
### get
292+
293+
Variable retrieval function that gets the value of a variable from the execution context. Commonly used to reference variables set via the `SET` statement in SQL.
294+
295+
**Function Signature**:
296+
297+
```java
298+
public static String eval(DataContext context, String key)
299+
```
300+
301+
**Parameter Description**:
302+
303+
| Parameter | Type | Description |
304+
|-----------|------|-------------|
305+
| `key` | String | Variable name |
306+
307+
**Return Value**: Returns the variable value (`String`), or `NULL` if the variable doesn't exist.
308+
309+
::: warning Note
310+
Since `get` is a SQL keyword, you need to wrap the function name with backticks when using it, written as `` `get` ``.
311+
:::
312+
313+
**Usage Example**:
314+
315+
```sql
316+
-- Set variable
317+
SET 'user_id' = '12345';
318+
319+
-- Get variable value
320+
SELECT `get`('user_id') AS user_id;
321+
322+
-- Use in expressions
323+
SELECT `get`('user_id') || '_suffix' AS user_id_with_suffix;
324+
325+
-- Type conversion
326+
SELECT CAST(`get`('limit_count') AS INT) AS limit_count;
327+
328+
-- Get variable name from table and use it
329+
CACHE TABLE var_names AS SELECT 'user_id' AS var_name;
330+
SELECT `get`(var_name) AS var_value FROM var_names;
331+
```
332+
333+
**Working Principle**:
334+
1. Function receives a variable name as parameter
335+
2. Looks up the corresponding variable value from the execution context (`ExecuteContext`)
336+
3. Returns the variable value, or `NULL` if the variable doesn't exist
337+
338+
**Typical Use Cases**:
339+
- Parameterized SQL queries
340+
- Dynamic configuration passing
341+
- Cross-statement variable sharing
342+
289343
## Custom UDF
290344

291345
Refer to [Programming Model](program_model.md#udf) documentation for how to develop custom UDFs.

‎sqlrec-common/src/main/java/com/sqlrec/common/runtime/SqlRecDataContext.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
package com.sqlrec.common.runtime;
2+
3+
import org.apache.calcite.DataContext;
4+
5+
public interface SqlRecDataContext extends DataContext {
6+
String getVariable(String key);
7+
}

‎sqlrec-core/src/main/java/com/sqlrec/runtime/CalciteBindable.java‎

Lines changed: 1 addition & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -3,24 +3,16 @@
33
import com.sqlrec.common.runtime.ExecuteContext;
44
import com.sqlrec.compiler.SqlTypeChecker;
55
import com.sqlrec.utils.NodeUtils;
6-
import org.apache.calcite.DataContext;
7-
import org.apache.calcite.adapter.java.JavaTypeFactory;
86
import org.apache.calcite.jdbc.CalciteSchema;
9-
import org.apache.calcite.jdbc.JavaTypeFactoryImpl;
107
import org.apache.calcite.linq4j.Enumerable;
118
import org.apache.calcite.linq4j.Linq4j;
12-
import org.apache.calcite.linq4j.QueryProvider;
139
import org.apache.calcite.rel.RelNode;
1410
import org.apache.calcite.rel.type.RelDataTypeField;
1511
import org.apache.calcite.runtime.Bindable;
16-
import org.apache.calcite.schema.SchemaPlus;
1712
import org.apache.calcite.sql.SqlNode;
18-
import org.checkerframework.checker.nullness.qual.Nullable;
1913

2014
import java.util.*;
2115

22-
import static org.apache.calcite.linq4j.Linq4j.DEFAULT_PROVIDER;
23-
2416
public class CalciteBindable extends BindableInterface {
2517
private Map<String, Object> parameters;
2618
private Bindable<Object[]> bindable;
@@ -59,7 +51,7 @@ public CalciteBindable(
5951

6052
@Override
6153
public Enumerable<Object[]> bind(CalciteSchema schema, ExecuteContext context) {
62-
Enumerable rawData = bindable.bind(new DataContextImpl(parameters, schema));
54+
Enumerable rawData = bindable.bind(new SqlRecDataContextImpl(parameters, schema, context));
6355

6456
List<Object[]> objArrayList = new ArrayList<>();
6557
for (Object obj : rawData) {
@@ -111,38 +103,4 @@ public String getPhysicalPlan() {
111103
public String getJavaExpression() {
112104
return javaExpression;
113105
}
114-
115-
public static class DataContextImpl implements DataContext {
116-
private Map<String, Object> parameters;
117-
private CalciteSchema schema;
118-
119-
public DataContextImpl(Map<String, Object> parameters, CalciteSchema schema) {
120-
this.parameters = parameters;
121-
this.schema = schema;
122-
}
123-
124-
@Override
125-
public @Nullable SchemaPlus getRootSchema() {
126-
return schema.plus();
127-
}
128-
129-
@Override
130-
public JavaTypeFactory getTypeFactory() {
131-
return new JavaTypeFactoryImpl();
132-
}
133-
134-
@Override
135-
public QueryProvider getQueryProvider() {
136-
return DEFAULT_PROVIDER;
137-
}
138-
139-
@Override
140-
public @Nullable Object get(String name) {
141-
switch (name) {
142-
case "currentTimestamp":
143-
return System.currentTimeMillis();
144-
}
145-
return parameters.get(name);
146-
}
147-
}
148106
}

‎sqlrec-core/src/main/java/com/sqlrec/runtime/SqlRecDataContextImpl.java‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
package com.sqlrec.runtime;
2+
3+
import com.sqlrec.common.runtime.ExecuteContext;
4+
import com.sqlrec.common.runtime.SqlRecDataContext;
5+
import org.apache.calcite.adapter.java.JavaTypeFactory;
6+
import org.apache.calcite.jdbc.CalciteSchema;
7+
import org.apache.calcite.jdbc.JavaTypeFactoryImpl;
8+
import org.apache.calcite.linq4j.QueryProvider;
9+
import org.apache.calcite.schema.SchemaPlus;
10+
import org.checkerframework.checker.nullness.qual.Nullable;
11+
12+
import java.util.Map;
13+
14+
import static org.apache.calcite.linq4j.Linq4j.DEFAULT_PROVIDER;
15+
16+
public class SqlRecDataContextImpl implements SqlRecDataContext {
17+
private Map<String, Object> parameters;
18+
private CalciteSchema schema;
19+
private ExecuteContext executeContext;
20+
21+
public SqlRecDataContextImpl(Map<String, Object> parameters, CalciteSchema schema, ExecuteContext executeContext) {
22+
this.parameters = parameters;
23+
this.schema = schema;
24+
this.executeContext = executeContext;
25+
}
26+
27+
@Override
28+
public @Nullable SchemaPlus getRootSchema() {
29+
return schema.plus();
30+
}
31+
32+
@Override
33+
public JavaTypeFactory getTypeFactory() {
34+
return new JavaTypeFactoryImpl();
35+
}
36+
37+
@Override
38+
public QueryProvider getQueryProvider() {
39+
return DEFAULT_PROVIDER;
40+
}
41+
42+
@Override
43+
public @Nullable Object get(String name) {
44+
switch (name) {
45+
case "currentTimestamp":
46+
return System.currentTimeMillis();
47+
}
48+
return parameters.get(name);
49+
}
50+
51+
@Override
52+
public String getVariable(String key) {
53+
return executeContext.getVariable(key);
54+
}
55+
}

‎sqlrec-core/src/main/java/com/sqlrec/schema/HmsSchema.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,9 @@
44
import com.google.common.collect.Multimap;
55
import com.sqlrec.common.config.SqlRecConfigs;
66
import com.sqlrec.common.utils.HiveTableUtils;
7+
import com.sqlrec.udf.UdfManager;
78
import com.sqlrec.udf.config.FunctionConfigs;
89
import com.sqlrec.utils.ObjCache;
9-
import com.sqlrec.utils.SchemaUtils;
1010
import org.apache.calcite.jdbc.CalciteSchema;
1111
import org.apache.calcite.schema.Function;
1212
import org.apache.calcite.schema.ScalarFunction;
@@ -123,13 +123,13 @@ private Multimap<String, Function> computeFunctionMap(Multimap<String, Function>
123123
List<String> functions = HmsClient.getAllFunctions(databaseName);
124124
for (String function : functions) {
125125
org.apache.hadoop.hive.metastore.api.Function functionObj = HmsClient.getFunctionObj(databaseName, function);
126-
ScalarFunction scalarFunction = SchemaUtils.createScalarFunction(functionObj.getClassName());
126+
ScalarFunction scalarFunction = UdfManager.createScalarFunction(functionObj.getClassName());
127127
if (scalarFunction != null) {
128128
functionMap.put(function, scalarFunction);
129129
}
130130
}
131131
for (Map.Entry<String, String> entry : FunctionConfigs.DEFAULT_SCALAR_FUNCTION_CONFIGS.entrySet()) {
132-
functionMap.put(entry.getKey(), SchemaUtils.createScalarFunction(entry.getValue()));
132+
functionMap.put(entry.getKey(), UdfManager.createScalarFunction(entry.getValue()));
133133
}
134134
} catch (Exception e) {
135135
log.error("Error while computing function map for schema {}", databaseName, e);

0 commit comments

Comments
 (0)