You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iotdb.apache.org by ja...@apache.org on 2026/06/27 01:43:57 UTC
(iotdb) branch master updated: Support IoTDBLocal in UDF for table model (#18035)
This is an automated email from the ASF dual-hosted git repository.
JackieTien97 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new fcd52d7ff06 Support IoTDBLocal in UDF for table model (#18035)
fcd52d7ff06 is described below
commit fcd52d7ff064e7644fd9429d9fa174c37103d0bf
Author: Weihao Li <60...@users.noreply.github.com>
AuthorDate: Sat Jun 27 09:43:51 2026 +0800
Support IoTDBLocal in UDF for table model (#18035)
---
.../iotdblocal/CreateTableProbeFunction.java | 30 +-
.../DeviceIdListInProcessTableFunction.java | 131 ++++++
.../iotdblocal/DeviceIdListTableFunction.java | 142 +++++++
.../DeviceNameEnrichBeforeStartTableFunction.java | 144 +++++++
.../DeviceNameEnrichInProcessTableFunction.java | 139 +++++++
.../relational/iotdblocal/DeviceNameFunction.java | 77 ++++
.../iotdblocal/DeviceSummaryFunction.java | 98 +++++
.../iotdblocal/DeviceSummaryNoCloseFunction.java | 93 +++++
.../iotdblocal/DropTableProbeFunction.java | 30 +-
.../relational/iotdblocal/InsertProbeFunction.java | 30 +-
.../iotdblocal/IoTDBLocalLogAggregateFunction.java | 133 ++++++
.../iotdblocal/IoTDBLocalLogFunction.java | 71 ++++
.../relational/iotdblocal/IoTDBLocalLogHelper.java | 93 +++++
.../iotdblocal/IoTDBLocalLogTableFunction.java | 134 ++++++
...lQueryUdafAtBeforeDestroyAggregateFunction.java | 101 +++++
...calQueryUdafAtBeforeStartAggregateFunction.java | 111 +++++
.../LocalQueryUdafInAddInputAggregateFunction.java | 123 ++++++
...alQueryUdafInCombineStateAggregateFunction.java | 123 ++++++
...calQueryUdafInOutputFinalAggregateFunction.java | 104 +++++
.../iotdblocal/MutatingQueryProbeFunction.java | 54 +++
.../iotdblocal/SecretTableQueryFunction.java | 57 +++
.../relational/it/db/it/udf/IoTDBLocalIT.java | 451 +++++++++++++++++++++
.../relational/it/db/it/udf/IoTDBLocalLogIT.java | 159 ++++++++
.../relational/it/db/it/udf/SQLFunctionUtils.java | 1 -
.../java/org/apache/iotdb/udf/api/IoTDBLocal.java | 70 ++++
...unctionLeafProcessor.java => UDFResultSet.java} | 38 +-
.../udf/api/relational/AggregateFunction.java | 33 ++
.../iotdb/udf/api/relational/ScalarFunction.java | 19 +
.../processor/TableFunctionDataProcessor.java | 34 ++
.../processor/TableFunctionLeafProcessor.java | 17 +
.../function/TableFunctionLeafOperator.java | 16 +-
.../process/function/TableFunctionOperator.java | 28 +-
.../relational/aggregation/AccumulatorFactory.java | 81 +++-
.../UserDefinedAggregateFunctionAccumulator.java | 65 ++-
.../GroupedUserDefinedAggregateAccumulator.java | 44 +-
.../relational/ColumnTransformerBuilder.java | 71 +++-
.../calc/plan/planner/TableOperatorGenerator.java | 111 ++++-
.../udf/UserDefineScalarFunctionTransformer.java | 46 ++-
.../iotdb/db/protocol/session/SessionManager.java | 25 ++
.../protocol/thrift/impl/ClientRPCServiceImpl.java | 4 +-
.../fragment/FragmentInstanceContext.java | 11 +
.../fragment/FragmentInstanceManager.java | 2 +
.../execution/schedule/task/DriverTask.java | 21 +-
.../iotdb/db/queryengine/plan/Coordinator.java | 42 ++
.../planner/DataNodeTableOperatorGenerator.java | 25 +-
.../plan/planner/LocalExecutionPlanContext.java | 4 +
.../db/queryengine/udf/InternalQueryExecutor.java | 133 ++++++
.../db/queryengine/udf/InternalQueryResult.java | 65 +++
.../iotdb/db/queryengine/udf/IoTDBLocalImpl.java | 147 +++++++
.../iotdb/db/queryengine/udf/UDFResultSetImpl.java | 154 +++++++
50 files changed, 3779 insertions(+), 156 deletions(-)
diff --git a/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionLeafProcessor.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/CreateTableProbeFunction.java
similarity index 51%
copy from iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionLeafProcessor.java
copy to integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/CreateTableProbeFunction.java
index ab8763f3c4a..90bbd4bac34 100644
--- a/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionLeafProcessor.java
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/CreateTableProbeFunction.java
@@ -17,31 +17,13 @@
* under the License.
*/
-package org.apache.iotdb.udf.api.relational.table.processor;
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
-import org.apache.tsfile.block.column.ColumnBuilder;
+/** Probes CREATE TABLE rejection for IoTDBLocal integration tests. */
+public class CreateTableProbeFunction extends MutatingQueryProbeFunction {
-import java.util.List;
-
-public interface TableFunctionLeafProcessor {
-
- default void beforeStart() {
- // do nothing
- }
-
- /**
- * This method processes a portion of data. It is called multiple times until the processor is
- * fully processed.
- *
- * @param columnBuilders a list of {@link ColumnBuilder} for each column in the output table.
- */
- void process(List<ColumnBuilder> columnBuilders);
-
- /** This method is called to determine if the processor has finished processing all data. */
- boolean isFinish();
-
- /** This method is mainly used to release the resources used in the UDF. */
- default void beforeDestroy() {
- // do nothing
+ @Override
+ protected String mutatingSql() {
+ return "CREATE TABLE should_not_exist_local_probe (x INT32 FIELD)";
}
}
diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceIdListInProcessTableFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceIdListInProcessTableFunction.java
new file mode 100644
index 00000000000..ae2a9cc4138
--- /dev/null
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceIdListInProcessTableFunction.java
@@ -0,0 +1,131 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
+
+import org.apache.iotdb.udf.api.IoTDBLocal;
+import org.apache.iotdb.udf.api.UDFResultSet;
+import org.apache.iotdb.udf.api.exception.UDFException;
+import org.apache.iotdb.udf.api.relational.TableFunction;
+import org.apache.iotdb.udf.api.relational.access.Record;
+import org.apache.iotdb.udf.api.relational.table.MapTableFunctionHandle;
+import org.apache.iotdb.udf.api.relational.table.TableFunctionAnalysis;
+import org.apache.iotdb.udf.api.relational.table.TableFunctionHandle;
+import org.apache.iotdb.udf.api.relational.table.TableFunctionProcessorProvider;
+import org.apache.iotdb.udf.api.relational.table.argument.Argument;
+import org.apache.iotdb.udf.api.relational.table.argument.DescribedSchema;
+import org.apache.iotdb.udf.api.relational.table.argument.ScalarArgument;
+import org.apache.iotdb.udf.api.relational.table.processor.TableFunctionLeafProcessor;
+import org.apache.iotdb.udf.api.relational.table.specification.ParameterSpecification;
+import org.apache.iotdb.udf.api.relational.table.specification.ScalarParameterSpecification;
+import org.apache.iotdb.udf.api.type.Type;
+
+import org.apache.tsfile.block.column.ColumnBuilder;
+import org.apache.tsfile.common.conf.TSFileConfig;
+import org.apache.tsfile.utils.Binary;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Leaf table function: loads device_id via {@link IoTDBLocal#query(String)} in {@link
+ * TableFunctionLeafProcessor#process(List, IoTDBLocal)}.
+ */
+public class DeviceIdListInProcessTableFunction implements TableFunction {
+
+ private static final String PREFIX_PARAMETER_NAME = "prefix";
+
+ @Override
+ public List<ParameterSpecification> getArgumentsSpecifications() {
+ return Collections.singletonList(
+ ScalarParameterSpecification.builder()
+ .name(PREFIX_PARAMETER_NAME)
+ .type(Type.STRING)
+ .defaultValue("")
+ .build());
+ }
+
+ @Override
+ public TableFunctionAnalysis analyze(Map<String, Argument> arguments) throws UDFException {
+ DescribedSchema schema = DescribedSchema.builder().addField("device_id", Type.STRING).build();
+ MapTableFunctionHandle handle =
+ new MapTableFunctionHandle.Builder()
+ .addProperty(
+ PREFIX_PARAMETER_NAME,
+ ((ScalarArgument) arguments.get(PREFIX_PARAMETER_NAME)).getValue())
+ .build();
+ return TableFunctionAnalysis.builder().properColumnSchema(schema).handle(handle).build();
+ }
+
+ @Override
+ public TableFunctionHandle createTableFunctionHandle() {
+ return new MapTableFunctionHandle();
+ }
+
+ @Override
+ public TableFunctionProcessorProvider getProcessorProvider(
+ TableFunctionHandle tableFunctionHandle) {
+ return new TableFunctionProcessorProvider() {
+ @Override
+ public TableFunctionLeafProcessor getSplitProcessor() {
+ return new InProcessProcessor(
+ (String)
+ ((MapTableFunctionHandle) tableFunctionHandle).getProperty(PREFIX_PARAMETER_NAME));
+ }
+ };
+ }
+
+ private static class InProcessProcessor implements TableFunctionLeafProcessor {
+ private final String prefix;
+ private boolean finish;
+
+ InProcessProcessor(String prefix) {
+ this.prefix = prefix == null ? "" : prefix;
+ }
+
+ @Override
+ public void process(List<ColumnBuilder> columnBuilders) {
+ throw new UnsupportedOperationException("framework invokes process(List, IoTDBLocal)");
+ }
+
+ @Override
+ public void process(List<ColumnBuilder> columnBuilders, IoTDBLocal local) {
+ if (finish) {
+ return;
+ }
+ try (UDFResultSet rs = local.query("SELECT device_id FROM device_info ORDER BY device_id")) {
+ while (rs.hasNext()) {
+ Record row = rs.next();
+ columnBuilders
+ .get(0)
+ .writeBinary(new Binary(prefix + row.getString(0), TSFileConfig.STRING_CHARSET));
+ }
+ } catch (UDFException e) {
+ throw new IllegalStateException(e);
+ }
+ finish = true;
+ }
+
+ @Override
+ public boolean isFinish() {
+ return finish;
+ }
+ }
+}
diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceIdListTableFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceIdListTableFunction.java
new file mode 100644
index 00000000000..1a036d4722a
--- /dev/null
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceIdListTableFunction.java
@@ -0,0 +1,142 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
+
+import org.apache.iotdb.udf.api.IoTDBLocal;
+import org.apache.iotdb.udf.api.UDFResultSet;
+import org.apache.iotdb.udf.api.exception.UDFException;
+import org.apache.iotdb.udf.api.relational.TableFunction;
+import org.apache.iotdb.udf.api.relational.access.Record;
+import org.apache.iotdb.udf.api.relational.table.MapTableFunctionHandle;
+import org.apache.iotdb.udf.api.relational.table.TableFunctionAnalysis;
+import org.apache.iotdb.udf.api.relational.table.TableFunctionHandle;
+import org.apache.iotdb.udf.api.relational.table.TableFunctionProcessorProvider;
+import org.apache.iotdb.udf.api.relational.table.argument.Argument;
+import org.apache.iotdb.udf.api.relational.table.argument.DescribedSchema;
+import org.apache.iotdb.udf.api.relational.table.argument.ScalarArgument;
+import org.apache.iotdb.udf.api.relational.table.processor.TableFunctionLeafProcessor;
+import org.apache.iotdb.udf.api.relational.table.specification.ParameterSpecification;
+import org.apache.iotdb.udf.api.relational.table.specification.ScalarParameterSpecification;
+import org.apache.iotdb.udf.api.type.Type;
+
+import org.apache.tsfile.block.column.ColumnBuilder;
+import org.apache.tsfile.common.conf.TSFileConfig;
+import org.apache.tsfile.utils.Binary;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Leaf table function: loads device_id via {@link IoTDBLocal#query(String)} in {@link
+ * TableFunctionLeafProcessor#beforeStart(IoTDBLocal)}.
+ */
+public class DeviceIdListTableFunction implements TableFunction {
+
+ private static final String PREFIX_PARAMETER_NAME = "prefix";
+
+ @Override
+ public List<ParameterSpecification> getArgumentsSpecifications() {
+ return Collections.singletonList(
+ ScalarParameterSpecification.builder()
+ .name(PREFIX_PARAMETER_NAME)
+ .type(Type.STRING)
+ .defaultValue("")
+ .build());
+ }
+
+ @Override
+ public TableFunctionAnalysis analyze(Map<String, Argument> arguments) throws UDFException {
+ DescribedSchema schema = DescribedSchema.builder().addField("device_id", Type.STRING).build();
+ MapTableFunctionHandle handle =
+ new MapTableFunctionHandle.Builder()
+ .addProperty(
+ PREFIX_PARAMETER_NAME,
+ ((ScalarArgument) arguments.get(PREFIX_PARAMETER_NAME)).getValue())
+ .build();
+ return TableFunctionAnalysis.builder().properColumnSchema(schema).handle(handle).build();
+ }
+
+ @Override
+ public TableFunctionHandle createTableFunctionHandle() {
+ return new MapTableFunctionHandle();
+ }
+
+ @Override
+ public TableFunctionProcessorProvider getProcessorProvider(
+ TableFunctionHandle tableFunctionHandle) {
+ return new TableFunctionProcessorProvider() {
+ @Override
+ public TableFunctionLeafProcessor getSplitProcessor() {
+ return new BeforeStartProcessor(
+ (String)
+ ((MapTableFunctionHandle) tableFunctionHandle).getProperty(PREFIX_PARAMETER_NAME));
+ }
+ };
+ }
+
+ private static class BeforeStartProcessor implements TableFunctionLeafProcessor {
+ private final String prefix;
+ private final List<String> deviceIds = new ArrayList<>();
+ private boolean finish;
+
+ BeforeStartProcessor(String prefix) {
+ this.prefix = prefix == null ? "" : prefix;
+ }
+
+ @Override
+ public void beforeStart(IoTDBLocal local) throws UDFException {
+ try (UDFResultSet rs = local.query("SELECT device_id FROM device_info ORDER BY device_id")) {
+ while (rs.hasNext()) {
+ Record row = rs.next();
+ deviceIds.add(row.getString(0));
+ }
+ }
+ }
+
+ @Override
+ public void process(List<ColumnBuilder> columnBuilders) {
+ writeDeviceIds(columnBuilders);
+ }
+
+ @Override
+ public void process(List<ColumnBuilder> columnBuilders, IoTDBLocal local) {
+ writeDeviceIds(columnBuilders);
+ }
+
+ private void writeDeviceIds(List<ColumnBuilder> columnBuilders) {
+ if (finish) {
+ return;
+ }
+ for (String deviceId : deviceIds) {
+ columnBuilders
+ .get(0)
+ .writeBinary(new Binary(prefix + deviceId, TSFileConfig.STRING_CHARSET));
+ }
+ finish = true;
+ }
+
+ @Override
+ public boolean isFinish() {
+ return finish;
+ }
+ }
+}
diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceNameEnrichBeforeStartTableFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceNameEnrichBeforeStartTableFunction.java
new file mode 100644
index 00000000000..af49b63d2f6
--- /dev/null
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceNameEnrichBeforeStartTableFunction.java
@@ -0,0 +1,144 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
+
+import org.apache.iotdb.udf.api.IoTDBLocal;
+import org.apache.iotdb.udf.api.UDFResultSet;
+import org.apache.iotdb.udf.api.exception.UDFArgumentNotValidException;
+import org.apache.iotdb.udf.api.exception.UDFException;
+import org.apache.iotdb.udf.api.relational.EmptyTableFunctionHandle;
+import org.apache.iotdb.udf.api.relational.TableFunction;
+import org.apache.iotdb.udf.api.relational.access.Record;
+import org.apache.iotdb.udf.api.relational.table.TableFunctionAnalysis;
+import org.apache.iotdb.udf.api.relational.table.TableFunctionHandle;
+import org.apache.iotdb.udf.api.relational.table.TableFunctionProcessorProvider;
+import org.apache.iotdb.udf.api.relational.table.argument.Argument;
+import org.apache.iotdb.udf.api.relational.table.argument.DescribedSchema;
+import org.apache.iotdb.udf.api.relational.table.argument.TableArgument;
+import org.apache.iotdb.udf.api.relational.table.processor.TableFunctionDataProcessor;
+import org.apache.iotdb.udf.api.relational.table.specification.ParameterSpecification;
+import org.apache.iotdb.udf.api.relational.table.specification.TableParameterSpecification;
+import org.apache.iotdb.udf.api.type.Type;
+
+import org.apache.tsfile.block.column.ColumnBuilder;
+import org.apache.tsfile.common.conf.TSFileConfig;
+import org.apache.tsfile.utils.Binary;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+/**
+ * Data table function: loads device_name mapping via {@link IoTDBLocal#query(String)} in {@link
+ * TableFunctionDataProcessor#beforeStart(IoTDBLocal)}.
+ */
+public class DeviceNameEnrichBeforeStartTableFunction implements TableFunction {
+
+ private static final String TABLE_PARAM = "DATA";
+
+ @Override
+ public List<ParameterSpecification> getArgumentsSpecifications() {
+ return Collections.singletonList(
+ TableParameterSpecification.builder().name(TABLE_PARAM).rowSemantics().build());
+ }
+
+ @Override
+ public TableFunctionAnalysis analyze(Map<String, Argument> arguments) throws UDFException {
+ TableArgument tableArgument = (TableArgument) arguments.get(TABLE_PARAM);
+ int deviceIdIndex = findDeviceIdIndex(tableArgument);
+ if (deviceIdIndex < 0) {
+ throw new UDFArgumentNotValidException("device_id column is required in table argument");
+ }
+ return TableFunctionAnalysis.builder()
+ .properColumnSchema(DescribedSchema.builder().addField("device_name", Type.STRING).build())
+ .requiredColumns(TABLE_PARAM, Collections.singletonList(deviceIdIndex))
+ .handle(new EmptyTableFunctionHandle())
+ .build();
+ }
+
+ @Override
+ public TableFunctionHandle createTableFunctionHandle() {
+ return new EmptyTableFunctionHandle();
+ }
+
+ @Override
+ public TableFunctionProcessorProvider getProcessorProvider(
+ TableFunctionHandle tableFunctionHandle) {
+ return new TableFunctionProcessorProvider() {
+ @Override
+ public TableFunctionDataProcessor getDataProcessor() {
+ return new BeforeStartProcessor();
+ }
+ };
+ }
+
+ private static int findDeviceIdIndex(TableArgument tableArgument) {
+ for (int i = 0; i < tableArgument.getFieldNames().size(); i++) {
+ Optional<String> fieldName = tableArgument.getFieldNames().get(i);
+ if (fieldName.isPresent() && "device_id".equalsIgnoreCase(fieldName.get())) {
+ return i;
+ }
+ }
+ return -1;
+ }
+
+ private static class BeforeStartProcessor implements TableFunctionDataProcessor {
+ private Map<String, String> idToName = Map.of();
+
+ @Override
+ public void beforeStart(IoTDBLocal local) throws UDFException {
+ Map<String, String> map = new HashMap<>();
+ try (UDFResultSet rs = local.query("SELECT device_id, device_name FROM device_info")) {
+ while (rs.hasNext()) {
+ Record row = rs.next();
+ map.put(row.getString(0), row.getString(1));
+ }
+ }
+ idToName = map;
+ }
+
+ @Override
+ public void process(
+ Record input,
+ List<ColumnBuilder> properColumnBuilders,
+ ColumnBuilder passThroughIndexBuilder) {
+ writeDeviceName(input, properColumnBuilders);
+ }
+
+ @Override
+ public void process(
+ Record input,
+ List<ColumnBuilder> properColumnBuilders,
+ ColumnBuilder passThroughIndexBuilder,
+ IoTDBLocal local) {
+ writeDeviceName(input, properColumnBuilders);
+ }
+
+ private void writeDeviceName(Record input, List<ColumnBuilder> properColumnBuilders) {
+ String deviceId = input.getString(0);
+ properColumnBuilders
+ .get(0)
+ .writeBinary(
+ new Binary(idToName.getOrDefault(deviceId, "未知设备"), TSFileConfig.STRING_CHARSET));
+ }
+ }
+}
diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceNameEnrichInProcessTableFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceNameEnrichInProcessTableFunction.java
new file mode 100644
index 00000000000..7839fa6b99a
--- /dev/null
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceNameEnrichInProcessTableFunction.java
@@ -0,0 +1,139 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
+
+import org.apache.iotdb.udf.api.IoTDBLocal;
+import org.apache.iotdb.udf.api.UDFResultSet;
+import org.apache.iotdb.udf.api.exception.UDFArgumentNotValidException;
+import org.apache.iotdb.udf.api.exception.UDFException;
+import org.apache.iotdb.udf.api.relational.EmptyTableFunctionHandle;
+import org.apache.iotdb.udf.api.relational.TableFunction;
+import org.apache.iotdb.udf.api.relational.access.Record;
+import org.apache.iotdb.udf.api.relational.table.TableFunctionAnalysis;
+import org.apache.iotdb.udf.api.relational.table.TableFunctionHandle;
+import org.apache.iotdb.udf.api.relational.table.TableFunctionProcessorProvider;
+import org.apache.iotdb.udf.api.relational.table.argument.Argument;
+import org.apache.iotdb.udf.api.relational.table.argument.DescribedSchema;
+import org.apache.iotdb.udf.api.relational.table.argument.TableArgument;
+import org.apache.iotdb.udf.api.relational.table.processor.TableFunctionDataProcessor;
+import org.apache.iotdb.udf.api.relational.table.specification.ParameterSpecification;
+import org.apache.iotdb.udf.api.relational.table.specification.TableParameterSpecification;
+import org.apache.iotdb.udf.api.type.Type;
+
+import org.apache.tsfile.block.column.ColumnBuilder;
+import org.apache.tsfile.common.conf.TSFileConfig;
+import org.apache.tsfile.utils.Binary;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+/**
+ * Data table function: looks up device_name via per-row {@link IoTDBLocal#query(String)} in {@link
+ * TableFunctionDataProcessor#process(Record, List, ColumnBuilder, IoTDBLocal)}.
+ */
+public class DeviceNameEnrichInProcessTableFunction implements TableFunction {
+
+ private static final String TABLE_PARAM = "DATA";
+
+ @Override
+ public List<ParameterSpecification> getArgumentsSpecifications() {
+ return Collections.singletonList(
+ TableParameterSpecification.builder().name(TABLE_PARAM).rowSemantics().build());
+ }
+
+ @Override
+ public TableFunctionAnalysis analyze(Map<String, Argument> arguments) throws UDFException {
+ TableArgument tableArgument = (TableArgument) arguments.get(TABLE_PARAM);
+ int deviceIdIndex = findDeviceIdIndex(tableArgument);
+ if (deviceIdIndex < 0) {
+ throw new UDFArgumentNotValidException("device_id column is required in table argument");
+ }
+ return TableFunctionAnalysis.builder()
+ .properColumnSchema(DescribedSchema.builder().addField("device_name", Type.STRING).build())
+ .requiredColumns(TABLE_PARAM, Collections.singletonList(deviceIdIndex))
+ .handle(new EmptyTableFunctionHandle())
+ .build();
+ }
+
+ @Override
+ public TableFunctionHandle createTableFunctionHandle() {
+ return new EmptyTableFunctionHandle();
+ }
+
+ @Override
+ public TableFunctionProcessorProvider getProcessorProvider(
+ TableFunctionHandle tableFunctionHandle) {
+ return new TableFunctionProcessorProvider() {
+ @Override
+ public TableFunctionDataProcessor getDataProcessor() {
+ return new InProcessProcessor();
+ }
+ };
+ }
+
+ private static int findDeviceIdIndex(TableArgument tableArgument) {
+ for (int i = 0; i < tableArgument.getFieldNames().size(); i++) {
+ Optional<String> fieldName = tableArgument.getFieldNames().get(i);
+ if (fieldName.isPresent() && "device_id".equalsIgnoreCase(fieldName.get())) {
+ return i;
+ }
+ }
+ return -1;
+ }
+
+ private static class InProcessProcessor implements TableFunctionDataProcessor {
+
+ @Override
+ public void process(
+ Record input,
+ List<ColumnBuilder> properColumnBuilders,
+ ColumnBuilder passThroughIndexBuilder) {
+ throw new UnsupportedOperationException(
+ "framework invokes process(Record, List, ColumnBuilder, IoTDBLocal)");
+ }
+
+ @Override
+ public void process(
+ Record input,
+ List<ColumnBuilder> properColumnBuilders,
+ ColumnBuilder passThroughIndexBuilder,
+ IoTDBLocal local) {
+ String deviceId = input.getString(0);
+ properColumnBuilders
+ .get(0)
+ .writeBinary(new Binary(queryDeviceName(local, deviceId), TSFileConfig.STRING_CHARSET));
+ }
+
+ private static String queryDeviceName(IoTDBLocal local, String deviceId) {
+ String sql =
+ String.format("SELECT device_name FROM device_info WHERE device_id = '%s'", deviceId);
+ try (UDFResultSet rs = local.query(sql)) {
+ if (rs.hasNext()) {
+ return rs.next().getString(0);
+ }
+ } catch (UDFException e) {
+ throw new IllegalStateException(e);
+ }
+ return "未知设备";
+ }
+ }
+}
diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceNameFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceNameFunction.java
new file mode 100644
index 00000000000..d08d0d0affe
--- /dev/null
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceNameFunction.java
@@ -0,0 +1,77 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
+
+import org.apache.iotdb.udf.api.IoTDBLocal;
+import org.apache.iotdb.udf.api.UDFResultSet;
+import org.apache.iotdb.udf.api.customizer.analysis.ScalarFunctionAnalysis;
+import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;
+import org.apache.iotdb.udf.api.exception.UDFArgumentNotValidException;
+import org.apache.iotdb.udf.api.exception.UDFException;
+import org.apache.iotdb.udf.api.relational.ScalarFunction;
+import org.apache.iotdb.udf.api.relational.access.Record;
+import org.apache.iotdb.udf.api.type.Type;
+
+import org.apache.tsfile.common.conf.TSFileConfig;
+import org.apache.tsfile.utils.Binary;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/** Lookup device_name by device_id via a single IoTDBLocal query in {@link #beforeStart}. */
+public class DeviceNameFunction implements ScalarFunction {
+
+ private Map<String, String> idToName = Map.of();
+
+ @Override
+ public ScalarFunctionAnalysis analyze(FunctionArguments arguments)
+ throws UDFArgumentNotValidException {
+ return new ScalarFunctionAnalysis.Builder().outputDataType(Type.STRING).build();
+ }
+
+ @Override
+ public void beforeStart(FunctionArguments arguments, IoTDBLocal local) throws UDFException {
+ local.info("DeviceNameFunction: loading device_info");
+ Map<String, String> map = new HashMap<>();
+ try (UDFResultSet rs = local.query("SELECT device_id, device_name FROM device_info")) {
+ while (rs.hasNext()) {
+ Record row = rs.next();
+ map.put(row.getString(0), row.getString(1));
+ }
+ }
+ idToName = map;
+ local.info("DeviceNameFunction: loaded {} mappings", idToName.size());
+ }
+
+ @Override
+ public Object evaluate(Record input) {
+ return lookupName(input);
+ }
+
+ @Override
+ public Object evaluate(Record input, IoTDBLocal local) {
+ return lookupName(input);
+ }
+
+ private Object lookupName(Record input) {
+ String deviceId = input.getString(0);
+ return new Binary(idToName.getOrDefault(deviceId, "未知设备"), TSFileConfig.STRING_CHARSET);
+ }
+}
diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceSummaryFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceSummaryFunction.java
new file mode 100644
index 00000000000..a400f70dcaf
--- /dev/null
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceSummaryFunction.java
@@ -0,0 +1,98 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
+
+import org.apache.iotdb.udf.api.IoTDBLocal;
+import org.apache.iotdb.udf.api.UDFResultSet;
+import org.apache.iotdb.udf.api.customizer.analysis.ScalarFunctionAnalysis;
+import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;
+import org.apache.iotdb.udf.api.exception.UDFArgumentNotValidException;
+import org.apache.iotdb.udf.api.exception.UDFException;
+import org.apache.iotdb.udf.api.relational.ScalarFunction;
+import org.apache.iotdb.udf.api.relational.access.Record;
+import org.apache.iotdb.udf.api.type.Type;
+
+import org.apache.tsfile.common.conf.TSFileConfig;
+import org.apache.tsfile.utils.Binary;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Lookup device_name and max_temp via two IoTDBLocal queries in {@link #beforeStart}, reading both
+ * {@link UDFResultSet} instances in an interleaved manner.
+ */
+public class DeviceSummaryFunction implements ScalarFunction {
+
+ private Map<String, String> idToName = Map.of();
+ private Map<String, Double> idToMaxTemp = Map.of();
+
+ @Override
+ public ScalarFunctionAnalysis analyze(FunctionArguments arguments)
+ throws UDFArgumentNotValidException {
+ return new ScalarFunctionAnalysis.Builder().outputDataType(Type.STRING).build();
+ }
+
+ @Override
+ public void beforeStart(FunctionArguments arguments, IoTDBLocal local) throws UDFException {
+ idToName = new HashMap<>();
+ idToMaxTemp = new HashMap<>();
+ UDFResultSet nameResult = local.query("SELECT device_id, device_name FROM device_info");
+ UDFResultSet limitResult = local.query("SELECT device_id, max_temp FROM device_limits");
+ try {
+ boolean hasName = nameResult.hasNext();
+ boolean hasLimit = limitResult.hasNext();
+ while (hasName || hasLimit) {
+ if (hasName) {
+ Record row = nameResult.next();
+ idToName.put(row.getString(0), row.getString(1));
+ hasName = nameResult.hasNext();
+ }
+ if (hasLimit) {
+ Record row = limitResult.next();
+ idToMaxTemp.put(row.getString(0), row.getDouble(1));
+ hasLimit = limitResult.hasNext();
+ }
+ }
+ } finally {
+ nameResult.close();
+ limitResult.close();
+ }
+ }
+
+ @Override
+ public Object evaluate(Record input) {
+ return buildSummary(input);
+ }
+
+ @Override
+ public Object evaluate(Record input, IoTDBLocal local) {
+ return buildSummary(input);
+ }
+
+ private Object buildSummary(Record input) {
+ String deviceId = input.getString(0);
+ String name = idToName.getOrDefault(deviceId, "未知设备");
+ Double maxTemp = idToMaxTemp.get(deviceId);
+ return new Binary(
+ String.format("%s(上限:%s)", name, maxTemp == null ? "未知" : maxTemp),
+ TSFileConfig.STRING_CHARSET);
+ }
+}
diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceSummaryNoCloseFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceSummaryNoCloseFunction.java
new file mode 100644
index 00000000000..c0ef353e4c1
--- /dev/null
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DeviceSummaryNoCloseFunction.java
@@ -0,0 +1,93 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
+
+import org.apache.iotdb.udf.api.IoTDBLocal;
+import org.apache.iotdb.udf.api.UDFResultSet;
+import org.apache.iotdb.udf.api.customizer.analysis.ScalarFunctionAnalysis;
+import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;
+import org.apache.iotdb.udf.api.exception.UDFArgumentNotValidException;
+import org.apache.iotdb.udf.api.exception.UDFException;
+import org.apache.iotdb.udf.api.relational.ScalarFunction;
+import org.apache.iotdb.udf.api.relational.access.Record;
+import org.apache.iotdb.udf.api.type.Type;
+
+import org.apache.tsfile.common.conf.TSFileConfig;
+import org.apache.tsfile.utils.Binary;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Same as {@link DeviceSummaryFunction} but intentionally does not close result sets, for verifying
+ * framework auto-cleanup.
+ */
+public class DeviceSummaryNoCloseFunction implements ScalarFunction {
+
+ private Map<String, String> idToName = Map.of();
+ private Map<String, Double> idToMaxTemp = Map.of();
+
+ @Override
+ public ScalarFunctionAnalysis analyze(FunctionArguments arguments)
+ throws UDFArgumentNotValidException {
+ return new ScalarFunctionAnalysis.Builder().outputDataType(Type.STRING).build();
+ }
+
+ @Override
+ public void beforeStart(FunctionArguments arguments, IoTDBLocal local) throws UDFException {
+ idToName = new HashMap<>();
+ idToMaxTemp = new HashMap<>();
+ UDFResultSet nameResult = local.query("SELECT device_id, device_name FROM device_info");
+ UDFResultSet limitResult = local.query("SELECT device_id, max_temp FROM device_limits");
+ boolean hasName = nameResult.hasNext();
+ boolean hasLimit = limitResult.hasNext();
+ while (hasName || hasLimit) {
+ if (hasName) {
+ Record row = nameResult.next();
+ idToName.put(row.getString(0), row.getString(1));
+ hasName = nameResult.hasNext();
+ }
+ if (hasLimit) {
+ Record row = limitResult.next();
+ idToMaxTemp.put(row.getString(0), row.getDouble(1));
+ hasLimit = limitResult.hasNext();
+ }
+ }
+ }
+
+ @Override
+ public Object evaluate(Record input) {
+ return buildSummary(input);
+ }
+
+ @Override
+ public Object evaluate(Record input, IoTDBLocal local) {
+ return buildSummary(input);
+ }
+
+ private Object buildSummary(Record input) {
+ String deviceId = input.getString(0);
+ String name = idToName.getOrDefault(deviceId, "未知设备");
+ Double maxTemp = idToMaxTemp.get(deviceId);
+ return new Binary(
+ String.format("%s(上限:%s)", name, maxTemp == null ? "未知" : maxTemp),
+ TSFileConfig.STRING_CHARSET);
+ }
+}
diff --git a/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionLeafProcessor.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DropTableProbeFunction.java
similarity index 51%
copy from iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionLeafProcessor.java
copy to integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DropTableProbeFunction.java
index ab8763f3c4a..a5825eab4fa 100644
--- a/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionLeafProcessor.java
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DropTableProbeFunction.java
@@ -17,31 +17,13 @@
* under the License.
*/
-package org.apache.iotdb.udf.api.relational.table.processor;
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
-import org.apache.tsfile.block.column.ColumnBuilder;
+/** Probes DROP TABLE rejection for IoTDBLocal integration tests. */
+public class DropTableProbeFunction extends MutatingQueryProbeFunction {
-import java.util.List;
-
-public interface TableFunctionLeafProcessor {
-
- default void beforeStart() {
- // do nothing
- }
-
- /**
- * This method processes a portion of data. It is called multiple times until the processor is
- * fully processed.
- *
- * @param columnBuilders a list of {@link ColumnBuilder} for each column in the output table.
- */
- void process(List<ColumnBuilder> columnBuilders);
-
- /** This method is called to determine if the processor has finished processing all data. */
- boolean isFinish();
-
- /** This method is mainly used to release the resources used in the UDF. */
- default void beforeDestroy() {
- // do nothing
+ @Override
+ protected String mutatingSql() {
+ return "DROP TABLE device_info";
}
}
diff --git a/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionLeafProcessor.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/InsertProbeFunction.java
similarity index 51%
copy from iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionLeafProcessor.java
copy to integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/InsertProbeFunction.java
index ab8763f3c4a..a4d5fc525c4 100644
--- a/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionLeafProcessor.java
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/InsertProbeFunction.java
@@ -17,31 +17,13 @@
* under the License.
*/
-package org.apache.iotdb.udf.api.relational.table.processor;
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
-import org.apache.tsfile.block.column.ColumnBuilder;
+/** Probes INSERT rejection for IoTDBLocal integration tests. */
+public class InsertProbeFunction extends MutatingQueryProbeFunction {
-import java.util.List;
-
-public interface TableFunctionLeafProcessor {
-
- default void beforeStart() {
- // do nothing
- }
-
- /**
- * This method processes a portion of data. It is called multiple times until the processor is
- * fully processed.
- *
- * @param columnBuilders a list of {@link ColumnBuilder} for each column in the output table.
- */
- void process(List<ColumnBuilder> columnBuilders);
-
- /** This method is called to determine if the processor has finished processing all data. */
- boolean isFinish();
-
- /** This method is mainly used to release the resources used in the UDF. */
- default void beforeDestroy() {
- // do nothing
+ @Override
+ protected String mutatingSql() {
+ return "INSERT INTO probe_table(time, device_id, value) VALUES (2, 'probe', 99.0)";
}
}
diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/IoTDBLocalLogAggregateFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/IoTDBLocalLogAggregateFunction.java
new file mode 100644
index 00000000000..7b7630bf5b0
--- /dev/null
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/IoTDBLocalLogAggregateFunction.java
@@ -0,0 +1,133 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
+
+import org.apache.iotdb.udf.api.IoTDBLocal;
+import org.apache.iotdb.udf.api.State;
+import org.apache.iotdb.udf.api.customizer.analysis.AggregateFunctionAnalysis;
+import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;
+import org.apache.iotdb.udf.api.exception.UDFArgumentNotValidException;
+import org.apache.iotdb.udf.api.relational.AggregateFunction;
+import org.apache.iotdb.udf.api.relational.access.Record;
+import org.apache.iotdb.udf.api.type.Type;
+import org.apache.iotdb.udf.api.utils.ResultValue;
+
+import java.nio.ByteBuffer;
+
+/** Exercises IoTDBLocal log APIs at each aggregate-function lifecycle hook. */
+public class IoTDBLocalLogAggregateFunction implements AggregateFunction {
+
+ private static class CountState implements State {
+ long count;
+ boolean addInputLogged;
+ boolean combineStateLogged;
+ boolean outputFinalLogged;
+ boolean destroyLogged;
+
+ @Override
+ public void reset() {
+ count = 0;
+ }
+
+ @Override
+ public byte[] serialize() {
+ ByteBuffer buffer = ByteBuffer.allocate(Long.BYTES);
+ buffer.putLong(count);
+ return buffer.array();
+ }
+
+ @Override
+ public void deserialize(byte[] bytes) {
+ count = ByteBuffer.wrap(bytes).getLong();
+ }
+ }
+
+ @Override
+ public AggregateFunctionAnalysis analyze(FunctionArguments arguments)
+ throws UDFArgumentNotValidException {
+ if (arguments.getArgumentsSize() != 1) {
+ throw new UDFArgumentNotValidException("IoTDBLocalLogAggregateFunction accepts one column");
+ }
+ return new AggregateFunctionAnalysis.Builder().outputDataType(Type.INT64).build();
+ }
+
+ @Override
+ public void beforeStart(FunctionArguments arguments, IoTDBLocal local) {
+ IoTDBLocalLogHelper.logAllApis(local, IoTDBLocalLogHelper.UDAF_BEFORE_START);
+ }
+
+ @Override
+ public State createState() {
+ return new CountState();
+ }
+
+ @Override
+ public void addInput(State state, Record input) {
+ CountState countState = (CountState) state;
+ countState.count++;
+ }
+
+ @Override
+ public void addInput(State state, Record input, IoTDBLocal local) {
+ CountState countState = (CountState) state;
+ if (!countState.addInputLogged) {
+ countState.addInputLogged = true;
+ IoTDBLocalLogHelper.logAllApis(local, IoTDBLocalLogHelper.UDAF_ADD_INPUT);
+ }
+ addInput(state, input);
+ }
+
+ @Override
+ public void combineState(State state, State rhs) {
+ CountState left = (CountState) state;
+ CountState right = (CountState) rhs;
+ left.count += right.count;
+ }
+
+ @Override
+ public void combineState(State state, State rhs, IoTDBLocal local) {
+ CountState countState = (CountState) state;
+ if (!countState.combineStateLogged) {
+ countState.combineStateLogged = true;
+ IoTDBLocalLogHelper.logAllApis(local, IoTDBLocalLogHelper.UDAF_COMBINE_STATE);
+ }
+ combineState(state, rhs);
+ }
+
+ @Override
+ public void outputFinal(State state, ResultValue resultValue) {
+ resultValue.setLong(((CountState) state).count);
+ }
+
+ @Override
+ public void outputFinal(State state, ResultValue resultValue, IoTDBLocal local) {
+ CountState countState = (CountState) state;
+ if (!countState.outputFinalLogged) {
+ countState.outputFinalLogged = true;
+ IoTDBLocalLogHelper.logAllApis(local, IoTDBLocalLogHelper.UDAF_OUTPUT_FINAL);
+ }
+ outputFinal(state, resultValue);
+ }
+
+ @Override
+ public void beforeDestroy(IoTDBLocal local) {
+ IoTDBLocalLogHelper.logAllApis(local, IoTDBLocalLogHelper.UDAF_BEFORE_DESTROY);
+ }
+}
diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/IoTDBLocalLogFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/IoTDBLocalLogFunction.java
new file mode 100644
index 00000000000..720828dcf1b
--- /dev/null
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/IoTDBLocalLogFunction.java
@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
+
+import org.apache.iotdb.udf.api.IoTDBLocal;
+import org.apache.iotdb.udf.api.customizer.analysis.ScalarFunctionAnalysis;
+import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;
+import org.apache.iotdb.udf.api.exception.UDFArgumentNotValidException;
+import org.apache.iotdb.udf.api.relational.ScalarFunction;
+import org.apache.iotdb.udf.api.relational.access.Record;
+import org.apache.iotdb.udf.api.type.Type;
+
+import org.apache.tsfile.common.conf.TSFileConfig;
+import org.apache.tsfile.utils.Binary;
+
+/** Exercises IoTDBLocal log APIs at each scalar-function lifecycle hook. */
+public class IoTDBLocalLogFunction implements ScalarFunction {
+
+ private boolean evaluateLogged;
+ private boolean destroyLogged;
+
+ @Override
+ public ScalarFunctionAnalysis analyze(FunctionArguments arguments)
+ throws UDFArgumentNotValidException {
+ return new ScalarFunctionAnalysis.Builder().outputDataType(Type.STRING).build();
+ }
+
+ @Override
+ public void beforeStart(FunctionArguments arguments, IoTDBLocal local) {
+ IoTDBLocalLogHelper.logAllApis(local, IoTDBLocalLogHelper.SCALAR_BEFORE_START);
+ }
+
+ @Override
+ public Object evaluate(Record input) {
+ return new Binary("ok", TSFileConfig.STRING_CHARSET);
+ }
+
+ @Override
+ public Object evaluate(Record input, IoTDBLocal local) {
+ if (!evaluateLogged) {
+ evaluateLogged = true;
+ IoTDBLocalLogHelper.logAllApis(local, IoTDBLocalLogHelper.SCALAR_EVALUATE);
+ }
+ return evaluate(input);
+ }
+
+ @Override
+ public void beforeDestroy(IoTDBLocal local) {
+ if (!destroyLogged) {
+ destroyLogged = true;
+ IoTDBLocalLogHelper.logAllApis(local, IoTDBLocalLogHelper.SCALAR_BEFORE_DESTROY);
+ }
+ }
+}
diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/IoTDBLocalLogHelper.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/IoTDBLocalLogHelper.java
new file mode 100644
index 00000000000..970b25c3e84
--- /dev/null
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/IoTDBLocalLogHelper.java
@@ -0,0 +1,93 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
+
+import org.apache.iotdb.udf.api.IoTDBLocal;
+
+/** Shared IoTDBLocal log markers for integration tests. */
+public final class IoTDBLocalLogHelper {
+
+ public static final String CAUSE_MESSAGE = "iotdb-local-it-log-cause";
+
+ public static final String SCALAR_BEFORE_START = "IOTDB_LOCAL_SCALAR_BEFORE_START";
+ public static final String SCALAR_EVALUATE = "IOTDB_LOCAL_SCALAR_EVALUATE";
+ public static final String SCALAR_BEFORE_DESTROY = "IOTDB_LOCAL_SCALAR_BEFORE_DESTROY";
+
+ public static final String UDAF_BEFORE_START = "IOTDB_LOCAL_UDAF_BEFORE_START";
+ public static final String UDAF_ADD_INPUT = "IOTDB_LOCAL_UDAF_ADD_INPUT";
+ public static final String UDAF_COMBINE_STATE = "IOTDB_LOCAL_UDAF_COMBINE_STATE";
+ public static final String UDAF_OUTPUT_FINAL = "IOTDB_LOCAL_UDAF_OUTPUT_FINAL";
+ public static final String UDAF_BEFORE_DESTROY = "IOTDB_LOCAL_UDAF_BEFORE_DESTROY";
+
+ public static final String TVF_BEFORE_START = "IOTDB_LOCAL_TVF_BEFORE_START";
+ public static final String TVF_PROCESS = "IOTDB_LOCAL_TVF_PROCESS";
+ public static final String TVF_BEFORE_DESTROY = "IOTDB_LOCAL_TVF_BEFORE_DESTROY";
+
+ private IoTDBLocalLogHelper() {}
+
+ public static void logAllApis(IoTDBLocal local, String phase) {
+ RuntimeException cause = new RuntimeException(CAUSE_MESSAGE);
+ local.info(phase + "_INFO_PLAIN");
+ local.info(phase + "_INFO_FORMAT loaded {} rows", 3);
+ local.info(phase + "_INFO_CAUSE", cause);
+ local.warn(phase + "_WARN_PLAIN");
+ local.warn(phase + "_WARN_FORMAT warn {} {}", "a", "b");
+ local.warn(phase + "_WARN_CAUSE", cause);
+ local.error(phase + "_ERROR_PLAIN");
+ local.error(phase + "_ERROR_FORMAT error code={}", 500);
+ local.error(phase + "_ERROR_CAUSE", cause);
+ }
+
+ public static String infoPlain(String phase) {
+ return phase + "_INFO_PLAIN";
+ }
+
+ public static String infoFormat(String phase) {
+ return phase + "_INFO_FORMAT loaded 3 rows";
+ }
+
+ public static String infoCause(String phase) {
+ return phase + "_INFO_CAUSE";
+ }
+
+ public static String warnPlain(String phase) {
+ return phase + "_WARN_PLAIN";
+ }
+
+ public static String warnFormat(String phase) {
+ return phase + "_WARN_FORMAT warn a b";
+ }
+
+ public static String warnCause(String phase) {
+ return phase + "_WARN_CAUSE";
+ }
+
+ public static String errorPlain(String phase) {
+ return phase + "_ERROR_PLAIN";
+ }
+
+ public static String errorFormat(String phase) {
+ return phase + "_ERROR_FORMAT error code=500";
+ }
+
+ public static String errorCause(String phase) {
+ return phase + "_ERROR_CAUSE";
+ }
+}
diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/IoTDBLocalLogTableFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/IoTDBLocalLogTableFunction.java
new file mode 100644
index 00000000000..39dfa057751
--- /dev/null
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/IoTDBLocalLogTableFunction.java
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
+
+import org.apache.iotdb.udf.api.IoTDBLocal;
+import org.apache.iotdb.udf.api.exception.UDFException;
+import org.apache.iotdb.udf.api.relational.TableFunction;
+import org.apache.iotdb.udf.api.relational.table.MapTableFunctionHandle;
+import org.apache.iotdb.udf.api.relational.table.TableFunctionAnalysis;
+import org.apache.iotdb.udf.api.relational.table.TableFunctionHandle;
+import org.apache.iotdb.udf.api.relational.table.TableFunctionProcessorProvider;
+import org.apache.iotdb.udf.api.relational.table.argument.Argument;
+import org.apache.iotdb.udf.api.relational.table.argument.DescribedSchema;
+import org.apache.iotdb.udf.api.relational.table.argument.ScalarArgument;
+import org.apache.iotdb.udf.api.relational.table.processor.TableFunctionLeafProcessor;
+import org.apache.iotdb.udf.api.relational.table.specification.ParameterSpecification;
+import org.apache.iotdb.udf.api.relational.table.specification.ScalarParameterSpecification;
+import org.apache.iotdb.udf.api.type.Type;
+
+import org.apache.tsfile.block.column.ColumnBuilder;
+import org.apache.tsfile.utils.Binary;
+
+import java.nio.charset.Charset;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+/** Exercises IoTDBLocal log APIs at each table-function lifecycle hook. */
+public class IoTDBLocalLogTableFunction implements TableFunction {
+
+ private static final String INPUT_PARAMETER_NAME = "input";
+
+ @Override
+ public List<ParameterSpecification> getArgumentsSpecifications() {
+ return Collections.singletonList(
+ ScalarParameterSpecification.builder()
+ .name(INPUT_PARAMETER_NAME)
+ .type(Type.STRING)
+ .build());
+ }
+
+ @Override
+ public TableFunctionAnalysis analyze(Map<String, Argument> arguments) throws UDFException {
+ DescribedSchema schema = DescribedSchema.builder().addField("output", Type.STRING).build();
+ MapTableFunctionHandle handle =
+ new MapTableFunctionHandle.Builder()
+ .addProperty(
+ INPUT_PARAMETER_NAME,
+ ((ScalarArgument) arguments.get(INPUT_PARAMETER_NAME)).getValue())
+ .build();
+ return TableFunctionAnalysis.builder().properColumnSchema(schema).handle(handle).build();
+ }
+
+ @Override
+ public TableFunctionHandle createTableFunctionHandle() {
+ return new MapTableFunctionHandle();
+ }
+
+ @Override
+ public TableFunctionProcessorProvider getProcessorProvider(
+ TableFunctionHandle tableFunctionHandle) {
+ return new TableFunctionProcessorProvider() {
+ @Override
+ public TableFunctionLeafProcessor getSplitProcessor() {
+ return new LogSplitProcessor(
+ (String)
+ ((MapTableFunctionHandle) tableFunctionHandle).getProperty(INPUT_PARAMETER_NAME));
+ }
+ };
+ }
+
+ private static class LogSplitProcessor implements TableFunctionLeafProcessor {
+ private final String input;
+ private boolean finish;
+ private boolean processLogged;
+ private boolean destroyLogged;
+
+ LogSplitProcessor(String input) {
+ this.input = input;
+ }
+
+ @Override
+ public void beforeStart(IoTDBLocal local) {
+ IoTDBLocalLogHelper.logAllApis(local, IoTDBLocalLogHelper.TVF_BEFORE_START);
+ }
+
+ @Override
+ public void process(List<ColumnBuilder> columnBuilders) {
+ for (String value : input.split(",")) {
+ columnBuilders.get(0).writeBinary(new Binary(value, Charset.defaultCharset()));
+ }
+ finish = true;
+ }
+
+ @Override
+ public void process(List<ColumnBuilder> columnBuilders, IoTDBLocal local) {
+ if (!processLogged) {
+ processLogged = true;
+ IoTDBLocalLogHelper.logAllApis(local, IoTDBLocalLogHelper.TVF_PROCESS);
+ }
+ process(columnBuilders);
+ }
+
+ @Override
+ public boolean isFinish() {
+ return finish;
+ }
+
+ @Override
+ public void beforeDestroy(IoTDBLocal local) {
+ if (!destroyLogged) {
+ destroyLogged = true;
+ IoTDBLocalLogHelper.logAllApis(local, IoTDBLocalLogHelper.TVF_BEFORE_DESTROY);
+ }
+ }
+ }
+}
diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/LocalQueryUdafAtBeforeDestroyAggregateFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/LocalQueryUdafAtBeforeDestroyAggregateFunction.java
new file mode 100644
index 00000000000..16e924c7bbd
--- /dev/null
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/LocalQueryUdafAtBeforeDestroyAggregateFunction.java
@@ -0,0 +1,101 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
+
+import org.apache.iotdb.udf.api.IoTDBLocal;
+import org.apache.iotdb.udf.api.State;
+import org.apache.iotdb.udf.api.UDFResultSet;
+import org.apache.iotdb.udf.api.customizer.analysis.AggregateFunctionAnalysis;
+import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;
+import org.apache.iotdb.udf.api.exception.UDFArgumentNotValidException;
+import org.apache.iotdb.udf.api.exception.UDFException;
+import org.apache.iotdb.udf.api.relational.AggregateFunction;
+import org.apache.iotdb.udf.api.relational.access.Record;
+import org.apache.iotdb.udf.api.type.Type;
+import org.apache.iotdb.udf.api.utils.ResultValue;
+
+import java.nio.ByteBuffer;
+
+/** UDAF that calls {@link IoTDBLocal#query(String)} only in {@link #beforeDestroy(IoTDBLocal)}. */
+public class LocalQueryUdafAtBeforeDestroyAggregateFunction implements AggregateFunction {
+
+ private static class CountState implements State {
+ long count;
+
+ @Override
+ public void reset() {
+ count = 0;
+ }
+
+ @Override
+ public byte[] serialize() {
+ ByteBuffer buffer = ByteBuffer.allocate(Long.BYTES);
+ buffer.putLong(count);
+ return buffer.array();
+ }
+
+ @Override
+ public void deserialize(byte[] bytes) {
+ count = ByteBuffer.wrap(bytes).getLong();
+ }
+ }
+
+ @Override
+ public AggregateFunctionAnalysis analyze(FunctionArguments arguments)
+ throws UDFArgumentNotValidException {
+ if (arguments.getArgumentsSize() != 1) {
+ throw new UDFArgumentNotValidException("expects one column");
+ }
+ return new AggregateFunctionAnalysis.Builder().outputDataType(Type.INT64).build();
+ }
+
+ @Override
+ public State createState() {
+ return new CountState();
+ }
+
+ @Override
+ public void addInput(State state, Record input) {
+ if (!input.isNull(0)) {
+ ((CountState) state).count++;
+ }
+ }
+
+ @Override
+ public void combineState(State state, State rhs) {
+ ((CountState) state).count += ((CountState) rhs).count;
+ }
+
+ @Override
+ public void outputFinal(State state, ResultValue resultValue) {
+ resultValue.setLong(((CountState) state).count);
+ }
+
+ @Override
+ public void beforeDestroy(IoTDBLocal local) {
+ try (UDFResultSet rs = local.query("SELECT COUNT(*) FROM device_limits")) {
+ while (rs.hasNext()) {
+ rs.next();
+ }
+ } catch (UDFException e) {
+ throw new IllegalStateException(e);
+ }
+ }
+}
diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/LocalQueryUdafAtBeforeStartAggregateFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/LocalQueryUdafAtBeforeStartAggregateFunction.java
new file mode 100644
index 00000000000..71dd5248514
--- /dev/null
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/LocalQueryUdafAtBeforeStartAggregateFunction.java
@@ -0,0 +1,111 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
+
+import org.apache.iotdb.udf.api.IoTDBLocal;
+import org.apache.iotdb.udf.api.State;
+import org.apache.iotdb.udf.api.UDFResultSet;
+import org.apache.iotdb.udf.api.customizer.analysis.AggregateFunctionAnalysis;
+import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;
+import org.apache.iotdb.udf.api.exception.UDFArgumentNotValidException;
+import org.apache.iotdb.udf.api.exception.UDFException;
+import org.apache.iotdb.udf.api.relational.AggregateFunction;
+import org.apache.iotdb.udf.api.relational.access.Record;
+import org.apache.iotdb.udf.api.type.Type;
+import org.apache.iotdb.udf.api.utils.ResultValue;
+
+import java.nio.ByteBuffer;
+
+/** UDAF that calls {@link IoTDBLocal#query(String)} only in {@link #beforeStart}. */
+public class LocalQueryUdafAtBeforeStartAggregateFunction implements AggregateFunction {
+
+ private long extraCount;
+
+ private static class CountState implements State {
+ long count;
+ long extraCount;
+
+ @Override
+ public void reset() {
+ count = 0;
+ }
+
+ @Override
+ public byte[] serialize() {
+ ByteBuffer buffer = ByteBuffer.allocate(Long.BYTES * 2);
+ buffer.putLong(count);
+ buffer.putLong(extraCount);
+ return buffer.array();
+ }
+
+ @Override
+ public void deserialize(byte[] bytes) {
+ ByteBuffer buffer = ByteBuffer.wrap(bytes);
+ count = buffer.getLong();
+ extraCount = buffer.getLong();
+ }
+ }
+
+ @Override
+ public AggregateFunctionAnalysis analyze(FunctionArguments arguments)
+ throws UDFArgumentNotValidException {
+ if (arguments.getArgumentsSize() != 1) {
+ throw new UDFArgumentNotValidException("expects one column");
+ }
+ return new AggregateFunctionAnalysis.Builder().outputDataType(Type.INT64).build();
+ }
+
+ @Override
+ public void beforeStart(FunctionArguments arguments, IoTDBLocal local) throws UDFException {
+ try (UDFResultSet rs = local.query("SELECT COUNT(*) FROM device_info")) {
+ if (rs.hasNext()) {
+ extraCount = rs.next().getLong(0);
+ }
+ }
+ }
+
+ @Override
+ public State createState() {
+ CountState state = new CountState();
+ state.extraCount = extraCount;
+ return state;
+ }
+
+ @Override
+ public void addInput(State state, Record input) {
+ if (!input.isNull(0)) {
+ ((CountState) state).count++;
+ }
+ }
+
+ @Override
+ public void combineState(State state, State rhs) {
+ CountState left = (CountState) state;
+ CountState right = (CountState) rhs;
+ left.count += right.count;
+ left.extraCount = Math.max(left.extraCount, right.extraCount);
+ }
+
+ @Override
+ public void outputFinal(State state, ResultValue resultValue) {
+ CountState countState = (CountState) state;
+ resultValue.setLong(countState.count + countState.extraCount);
+ }
+}
diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/LocalQueryUdafInAddInputAggregateFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/LocalQueryUdafInAddInputAggregateFunction.java
new file mode 100644
index 00000000000..1f1005925cd
--- /dev/null
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/LocalQueryUdafInAddInputAggregateFunction.java
@@ -0,0 +1,123 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
+
+import org.apache.iotdb.udf.api.IoTDBLocal;
+import org.apache.iotdb.udf.api.State;
+import org.apache.iotdb.udf.api.UDFResultSet;
+import org.apache.iotdb.udf.api.customizer.analysis.AggregateFunctionAnalysis;
+import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;
+import org.apache.iotdb.udf.api.exception.UDFArgumentNotValidException;
+import org.apache.iotdb.udf.api.exception.UDFException;
+import org.apache.iotdb.udf.api.relational.AggregateFunction;
+import org.apache.iotdb.udf.api.relational.access.Record;
+import org.apache.iotdb.udf.api.type.Type;
+import org.apache.iotdb.udf.api.utils.ResultValue;
+
+import java.nio.ByteBuffer;
+
+/**
+ * UDAF that calls {@link IoTDBLocal#query(String)} only in {@link #addInput(State, Record,
+ * IoTDBLocal)}.
+ */
+public class LocalQueryUdafInAddInputAggregateFunction implements AggregateFunction {
+
+ private static class CountState implements State {
+ long count;
+ long extraCount;
+ boolean queried;
+
+ @Override
+ public void reset() {
+ count = 0;
+ extraCount = 0;
+ queried = false;
+ }
+
+ @Override
+ public byte[] serialize() {
+ ByteBuffer buffer = ByteBuffer.allocate(Long.BYTES * 2 + 1);
+ buffer.putLong(count);
+ buffer.putLong(extraCount);
+ buffer.put((byte) (queried ? 1 : 0));
+ return buffer.array();
+ }
+
+ @Override
+ public void deserialize(byte[] bytes) {
+ ByteBuffer buffer = ByteBuffer.wrap(bytes);
+ count = buffer.getLong();
+ extraCount = buffer.getLong();
+ queried = buffer.get() != 0;
+ }
+ }
+
+ @Override
+ public AggregateFunctionAnalysis analyze(FunctionArguments arguments)
+ throws UDFArgumentNotValidException {
+ if (arguments.getArgumentsSize() != 1) {
+ throw new UDFArgumentNotValidException("expects one column");
+ }
+ return new AggregateFunctionAnalysis.Builder().outputDataType(Type.INT64).build();
+ }
+
+ @Override
+ public State createState() {
+ return new CountState();
+ }
+
+ @Override
+ public void addInput(State state, Record input) {
+ if (!input.isNull(0)) {
+ ((CountState) state).count++;
+ }
+ }
+
+ @Override
+ public void addInput(State state, Record input, IoTDBLocal local) {
+ CountState countState = (CountState) state;
+ if (!countState.queried) {
+ try (UDFResultSet rs = local.query("SELECT COUNT(*) FROM device_limits")) {
+ if (rs.hasNext()) {
+ countState.extraCount = rs.next().getLong(0);
+ }
+ } catch (UDFException e) {
+ throw new IllegalStateException(e);
+ }
+ countState.queried = true;
+ }
+ addInput(state, input);
+ }
+
+ @Override
+ public void combineState(State state, State rhs) {
+ CountState left = (CountState) state;
+ CountState right = (CountState) rhs;
+ left.count += right.count;
+ left.extraCount = Math.max(left.extraCount, right.extraCount);
+ left.queried = left.queried || right.queried;
+ }
+
+ @Override
+ public void outputFinal(State state, ResultValue resultValue) {
+ CountState countState = (CountState) state;
+ resultValue.setLong(countState.count + countState.extraCount);
+ }
+}
diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/LocalQueryUdafInCombineStateAggregateFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/LocalQueryUdafInCombineStateAggregateFunction.java
new file mode 100644
index 00000000000..7e66de05bc2
--- /dev/null
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/LocalQueryUdafInCombineStateAggregateFunction.java
@@ -0,0 +1,123 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
+
+import org.apache.iotdb.udf.api.IoTDBLocal;
+import org.apache.iotdb.udf.api.State;
+import org.apache.iotdb.udf.api.UDFResultSet;
+import org.apache.iotdb.udf.api.customizer.analysis.AggregateFunctionAnalysis;
+import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;
+import org.apache.iotdb.udf.api.exception.UDFArgumentNotValidException;
+import org.apache.iotdb.udf.api.exception.UDFException;
+import org.apache.iotdb.udf.api.relational.AggregateFunction;
+import org.apache.iotdb.udf.api.relational.access.Record;
+import org.apache.iotdb.udf.api.type.Type;
+import org.apache.iotdb.udf.api.utils.ResultValue;
+
+import java.nio.ByteBuffer;
+
+/**
+ * UDAF that calls {@link IoTDBLocal#query(String)} only in {@link #combineState(State, State,
+ * IoTDBLocal)}.
+ */
+public class LocalQueryUdafInCombineStateAggregateFunction implements AggregateFunction {
+
+ private static class CountState implements State {
+ long count;
+ long extraCount;
+ boolean combineQueried;
+
+ @Override
+ public void reset() {
+ count = 0;
+ extraCount = 0;
+ combineQueried = false;
+ }
+
+ @Override
+ public byte[] serialize() {
+ ByteBuffer buffer = ByteBuffer.allocate(Long.BYTES * 2 + 1);
+ buffer.putLong(count);
+ buffer.putLong(extraCount);
+ buffer.put((byte) (combineQueried ? 1 : 0));
+ return buffer.array();
+ }
+
+ @Override
+ public void deserialize(byte[] bytes) {
+ ByteBuffer buffer = ByteBuffer.wrap(bytes);
+ count = buffer.getLong();
+ extraCount = buffer.getLong();
+ combineQueried = buffer.get() != 0;
+ }
+ }
+
+ @Override
+ public AggregateFunctionAnalysis analyze(FunctionArguments arguments)
+ throws UDFArgumentNotValidException {
+ if (arguments.getArgumentsSize() != 1) {
+ throw new UDFArgumentNotValidException("expects one column");
+ }
+ return new AggregateFunctionAnalysis.Builder().outputDataType(Type.INT64).build();
+ }
+
+ @Override
+ public State createState() {
+ return new CountState();
+ }
+
+ @Override
+ public void addInput(State state, Record input) {
+ if (!input.isNull(0)) {
+ ((CountState) state).count++;
+ }
+ }
+
+ @Override
+ public void combineState(State state, State rhs) {
+ CountState left = (CountState) state;
+ CountState right = (CountState) rhs;
+ left.count += right.count;
+ left.extraCount = Math.max(left.extraCount, right.extraCount);
+ left.combineQueried = left.combineQueried || right.combineQueried;
+ }
+
+ @Override
+ public void combineState(State state, State rhs, IoTDBLocal local) {
+ CountState left = (CountState) state;
+ if (!left.combineQueried) {
+ try (UDFResultSet rs = local.query("SELECT COUNT(*) FROM device_limits")) {
+ if (rs.hasNext()) {
+ left.extraCount = rs.next().getLong(0);
+ }
+ } catch (UDFException e) {
+ throw new IllegalStateException(e);
+ }
+ left.combineQueried = true;
+ }
+ combineState(state, rhs);
+ }
+
+ @Override
+ public void outputFinal(State state, ResultValue resultValue) {
+ CountState countState = (CountState) state;
+ resultValue.setLong(countState.count + countState.extraCount);
+ }
+}
diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/LocalQueryUdafInOutputFinalAggregateFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/LocalQueryUdafInOutputFinalAggregateFunction.java
new file mode 100644
index 00000000000..bf2dbaae526
--- /dev/null
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/LocalQueryUdafInOutputFinalAggregateFunction.java
@@ -0,0 +1,104 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
+
+import org.apache.iotdb.udf.api.IoTDBLocal;
+import org.apache.iotdb.udf.api.State;
+import org.apache.iotdb.udf.api.UDFResultSet;
+import org.apache.iotdb.udf.api.customizer.analysis.AggregateFunctionAnalysis;
+import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;
+import org.apache.iotdb.udf.api.exception.UDFArgumentNotValidException;
+import org.apache.iotdb.udf.api.exception.UDFException;
+import org.apache.iotdb.udf.api.relational.AggregateFunction;
+import org.apache.iotdb.udf.api.relational.access.Record;
+import org.apache.iotdb.udf.api.type.Type;
+import org.apache.iotdb.udf.api.utils.ResultValue;
+
+import java.nio.ByteBuffer;
+
+/** UDAF that calls {@link IoTDBLocal#query(String)} only in {@link #outputFinal} with local. */
+public class LocalQueryUdafInOutputFinalAggregateFunction implements AggregateFunction {
+
+ private static class CountState implements State {
+ long count;
+
+ @Override
+ public void reset() {
+ count = 0;
+ }
+
+ @Override
+ public byte[] serialize() {
+ ByteBuffer buffer = ByteBuffer.allocate(Long.BYTES);
+ buffer.putLong(count);
+ return buffer.array();
+ }
+
+ @Override
+ public void deserialize(byte[] bytes) {
+ count = ByteBuffer.wrap(bytes).getLong();
+ }
+ }
+
+ @Override
+ public AggregateFunctionAnalysis analyze(FunctionArguments arguments)
+ throws UDFArgumentNotValidException {
+ if (arguments.getArgumentsSize() != 1) {
+ throw new UDFArgumentNotValidException("expects one column");
+ }
+ return new AggregateFunctionAnalysis.Builder().outputDataType(Type.INT64).build();
+ }
+
+ @Override
+ public State createState() {
+ return new CountState();
+ }
+
+ @Override
+ public void addInput(State state, Record input) {
+ if (!input.isNull(0)) {
+ ((CountState) state).count++;
+ }
+ }
+
+ @Override
+ public void combineState(State state, State rhs) {
+ ((CountState) state).count += ((CountState) rhs).count;
+ }
+
+ @Override
+ public void outputFinal(State state, ResultValue resultValue) {
+ resultValue.setLong(((CountState) state).count);
+ }
+
+ @Override
+ public void outputFinal(State state, ResultValue resultValue, IoTDBLocal local) {
+ CountState countState = (CountState) state;
+ long extraCount = 0;
+ try (UDFResultSet rs = local.query("SELECT COUNT(*) FROM device_info")) {
+ if (rs.hasNext()) {
+ extraCount = rs.next().getLong(0);
+ }
+ } catch (UDFException e) {
+ throw new IllegalStateException(e);
+ }
+ resultValue.setLong(countState.count + extraCount);
+ }
+}
diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/MutatingQueryProbeFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/MutatingQueryProbeFunction.java
new file mode 100644
index 00000000000..09c1784800e
--- /dev/null
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/MutatingQueryProbeFunction.java
@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
+
+import org.apache.iotdb.udf.api.IoTDBLocal;
+import org.apache.iotdb.udf.api.customizer.analysis.ScalarFunctionAnalysis;
+import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;
+import org.apache.iotdb.udf.api.exception.UDFArgumentNotValidException;
+import org.apache.iotdb.udf.api.exception.UDFException;
+import org.apache.iotdb.udf.api.relational.ScalarFunction;
+import org.apache.iotdb.udf.api.relational.access.Record;
+import org.apache.iotdb.udf.api.type.Type;
+
+import org.apache.tsfile.common.conf.TSFileConfig;
+import org.apache.tsfile.utils.Binary;
+
+/** Base scalar UDF that probes a mutating SQL statement via {@link IoTDBLocal#query(String)}. */
+abstract class MutatingQueryProbeFunction implements ScalarFunction {
+
+ protected abstract String mutatingSql();
+
+ @Override
+ public ScalarFunctionAnalysis analyze(FunctionArguments arguments)
+ throws UDFArgumentNotValidException {
+ return new ScalarFunctionAnalysis.Builder().outputDataType(Type.STRING).build();
+ }
+
+ @Override
+ public void beforeStart(FunctionArguments arguments, IoTDBLocal local) throws UDFException {
+ local.query(mutatingSql());
+ }
+
+ @Override
+ public Object evaluate(Record input) {
+ return new Binary("unexpected", TSFileConfig.STRING_CHARSET);
+ }
+}
diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/SecretTableQueryFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/SecretTableQueryFunction.java
new file mode 100644
index 00000000000..de3b8a13aae
--- /dev/null
+++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/SecretTableQueryFunction.java
@@ -0,0 +1,57 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.query.udf.example.relational.iotdblocal;
+
+import org.apache.iotdb.udf.api.IoTDBLocal;
+import org.apache.iotdb.udf.api.UDFResultSet;
+import org.apache.iotdb.udf.api.customizer.analysis.ScalarFunctionAnalysis;
+import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;
+import org.apache.iotdb.udf.api.exception.UDFArgumentNotValidException;
+import org.apache.iotdb.udf.api.exception.UDFException;
+import org.apache.iotdb.udf.api.relational.ScalarFunction;
+import org.apache.iotdb.udf.api.relational.access.Record;
+import org.apache.iotdb.udf.api.type.Type;
+
+import org.apache.tsfile.common.conf.TSFileConfig;
+import org.apache.tsfile.utils.Binary;
+
+/** Probes IoTDBLocal query against secret_table for permission integration tests. */
+public class SecretTableQueryFunction implements ScalarFunction {
+
+ @Override
+ public ScalarFunctionAnalysis analyze(FunctionArguments arguments)
+ throws UDFArgumentNotValidException {
+ return new ScalarFunctionAnalysis.Builder().outputDataType(Type.STRING).build();
+ }
+
+ @Override
+ public void beforeStart(FunctionArguments arguments, IoTDBLocal local) throws UDFException {
+ try (UDFResultSet rs = local.query("SELECT secret FROM secret_table")) {
+ while (rs.hasNext()) {
+ rs.next();
+ }
+ }
+ }
+
+ @Override
+ public Object evaluate(Record input) {
+ return new Binary(input.getString(0), TSFileConfig.STRING_CHARSET);
+ }
+}
diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBLocalIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBLocalIT.java
new file mode 100644
index 00000000000..4e918e8bbef
--- /dev/null
+++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBLocalIT.java
@@ -0,0 +1,451 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.relational.it.db.it.udf;
+
+import org.apache.iotdb.it.env.EnvFactory;
+import org.apache.iotdb.it.framework.IoTDBTestRunner;
+import org.apache.iotdb.itbase.category.TableClusterIT;
+import org.apache.iotdb.itbase.category.TableLocalStandaloneIT;
+
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+
+import java.sql.Connection;
+import java.sql.Statement;
+
+import static org.apache.iotdb.db.it.utils.TestUtils.tableAssertTestFail;
+import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest;
+import static org.junit.Assert.fail;
+import static org.junit.Assume.assumeTrue;
+
+/**
+ * Integration tests for {@link org.apache.iotdb.udf.api.IoTDBLocal} in table-model UDF, covering
+ * compatibility, embedded query, permission inheritance and auto resource cleanup.
+ */
+@RunWith(IoTDBTestRunner.class)
+@Category({TableLocalStandaloneIT.class, TableClusterIT.class})
+public class IoTDBLocalIT {
+
+ private static final String DATABASE_NAME = "iotdb_local_it";
+ private static final String LIMITED_USER = "iotdb_local_user";
+ private static final String LIMITED_PASSWORD = "iotdbLocalPw123456";
+
+ private static final String PKG = "org.apache.iotdb.db.query.udf.example.relational";
+ private static final String IOTDB_LOCAL_PKG =
+ "org.apache.iotdb.db.query.udf.example.relational.iotdblocal";
+
+ private static final String[] SETUP_SQLS =
+ new String[] {
+ "CREATE DATABASE " + DATABASE_NAME,
+ "USE " + DATABASE_NAME,
+ "CREATE TABLE readings (device_id STRING TAG, temperature DOUBLE FIELD)",
+ "CREATE TABLE device_info (device_id STRING TAG, device_name STRING FIELD)",
+ "CREATE TABLE device_limits (device_id STRING TAG, max_temp DOUBLE FIELD)",
+ "CREATE TABLE secret_table (device_id STRING TAG, secret STRING FIELD)",
+ "CREATE TABLE vehicle (device_id STRING TAG, s1 INT32 FIELD, s2 INT64 FIELD)",
+ "CREATE TABLE probe_table (device_id STRING TAG, value DOUBLE FIELD)",
+ "INSERT INTO device_info(time, device_id, device_name) VALUES (1, 'd1', '一号车间温度传感器'), (1, 'd2', '二号车间温度传感器')",
+ "INSERT INTO device_limits(time, device_id, max_temp) VALUES (1, 'd1', 30.0), (1, 'd2', 35.0)",
+ "INSERT INTO readings(time, device_id, temperature) VALUES (1000, 'd1', 25.5), (1001, 'd2', 32.0), (1002, 'd3', 20.0)",
+ "INSERT INTO secret_table(time, device_id, secret) VALUES (1, 'd1', 'top-secret')",
+ "INSERT INTO probe_table(time, device_id, value) VALUES (1, 'seed', 1.0)",
+ "INSERT INTO vehicle(time, device_id, s1, s2) VALUES (1, 'd0', 1, 1)",
+ "INSERT INTO vehicle(time, device_id, s1, s2) VALUES (2, 'd0', null, 2)",
+ "INSERT INTO vehicle(time, device_id, s1, s2) VALUES (3, 'd0', 3, 3)",
+ "INSERT INTO vehicle(time, device_id, s1) VALUES (5, 'd0', 4)",
+ "FLUSH",
+ "CLEAR ATTRIBUTE CACHE",
+ };
+
+ @BeforeClass
+ public static void setUp() throws Exception {
+ EnvFactory.getEnv().getConfig().getCommonConfig().setEnforceStrongPassword(false);
+ EnvFactory.getEnv().initClusterEnvironment();
+ executeAsRoot(SETUP_SQLS);
+ }
+
+ @AfterClass
+ public static void tearDown() throws Exception {
+ EnvFactory.getEnv().cleanClusterEnvironment();
+ }
+
+ @After
+ public void dropFunctions() {
+ SQLFunctionUtils.dropAllUDF();
+ }
+
+ private static void executeAsRoot(String... sqls) {
+ try (Connection connection = EnvFactory.getEnv().getTableConnection();
+ Statement statement = connection.createStatement()) {
+ for (String sql : sqls) {
+ statement.execute(sql);
+ }
+ } catch (Exception e) {
+ fail("executeAsRoot failed: " + e.getMessage());
+ }
+ }
+
+ // ── compatibility: legacy UDF without IoTDBLocal parameter ─────────────────
+
+ @Test
+ public void testLegacyScalarUdf() {
+ SQLFunctionUtils.createUDF("contain_null", PKG + ".ContainNull");
+ SQLFunctionUtils.createUDF("all_sum", PKG + ".AllSum");
+ tableResultSetEqualTest(
+ "SELECT time, contain_null(s1, s2) AS contain_null, contain_null(s1) AS s1_null FROM vehicle",
+ new String[] {"time", "contain_null", "s1_null"},
+ new String[] {
+ "1970-01-01T00:00:00.001Z,false,false,",
+ "1970-01-01T00:00:00.002Z,true,true,",
+ "1970-01-01T00:00:00.003Z,false,false,",
+ "1970-01-01T00:00:00.005Z,true,false,"
+ },
+ DATABASE_NAME);
+ tableResultSetEqualTest(
+ "SELECT time, all_sum(s1, s2) AS s12 FROM vehicle",
+ new String[] {"time", "s12"},
+ new String[] {
+ "1970-01-01T00:00:00.001Z,2,",
+ "1970-01-01T00:00:00.002Z,2,",
+ "1970-01-01T00:00:00.003Z,6,",
+ "1970-01-01T00:00:00.005Z,4,"
+ },
+ DATABASE_NAME);
+ }
+
+ @Test
+ public void testLegacyAggregateUdf() {
+ SQLFunctionUtils.createUDF("my_avg", PKG + ".MyAvg");
+ tableResultSetEqualTest(
+ "SELECT device_id, my_avg(s1) AS avg_s1 FROM vehicle GROUP BY device_id",
+ new String[] {"device_id", "avg_s1"},
+ new String[] {"d0,2.6666666666666665,"},
+ DATABASE_NAME);
+ }
+
+ @Test
+ public void testLegacyTableFunctionUdf() {
+ SQLFunctionUtils.createUDF("my_split", PKG + ".MySplit");
+ SQLFunctionUtils.createUDF("my_repeat", PKG + ".MyRepeatWithoutIndex");
+ tableResultSetEqualTest(
+ "select * from my_split('a,b,c')",
+ new String[] {"output"},
+ new String[] {"a,", "b,", "c,"},
+ DATABASE_NAME);
+ tableResultSetEqualTest(
+ "select * from my_repeat((select time, device_id, s1 from vehicle), 2) order by time",
+ new String[] {"time", "device_id", "s1"},
+ new String[] {
+ "1970-01-01T00:00:00.001Z,d0,1,",
+ "1970-01-01T00:00:00.001Z,d0,1,",
+ "1970-01-01T00:00:00.002Z,d0,null,",
+ "1970-01-01T00:00:00.002Z,d0,null,",
+ "1970-01-01T00:00:00.003Z,d0,3,",
+ "1970-01-01T00:00:00.003Z,d0,3,",
+ "1970-01-01T00:00:00.005Z,d0,4,",
+ "1970-01-01T00:00:00.005Z,d0,4,",
+ },
+ DATABASE_NAME);
+ }
+
+ // ── IoTDBLocal embedded query ───────────────────────────────────────────────
+
+ @Test
+ public void testDeviceNameWithSingleQuery() {
+ SQLFunctionUtils.createUDF("device_name", IOTDB_LOCAL_PKG + ".DeviceNameFunction");
+ tableResultSetEqualTest(
+ "SELECT time, device_id, device_name(device_id) AS name, temperature FROM readings ORDER BY time",
+ new String[] {"time", "device_id", "name", "temperature"},
+ new String[] {
+ "1970-01-01T00:00:01.000Z,d1,一号车间温度传感器,25.5,",
+ "1970-01-01T00:00:01.001Z,d2,二号车间温度传感器,32.0,",
+ "1970-01-01T00:00:01.002Z,d3,未知设备,20.0,",
+ },
+ DATABASE_NAME);
+ }
+
+ @Test
+ public void testInterleavedIoTDBLocalQueries() {
+ SQLFunctionUtils.createUDF("device_summary", IOTDB_LOCAL_PKG + ".DeviceSummaryFunction");
+ tableResultSetEqualTest(
+ "SELECT time, device_id, temperature, device_summary(device_id) AS summary FROM readings ORDER BY time",
+ new String[] {"time", "device_id", "temperature", "summary"},
+ new String[] {
+ "1970-01-01T00:00:01.000Z,d1,25.5,一号车间温度传感器(上限:30.0),",
+ "1970-01-01T00:00:01.001Z,d2,32.0,二号车间温度传感器(上限:35.0),",
+ "1970-01-01T00:00:01.002Z,d3,20.0,未知设备(上限:未知),",
+ },
+ DATABASE_NAME);
+ }
+
+ @Test
+ public void testDeviceSummaryWithoutManualClose() {
+ SQLFunctionUtils.createUDF(
+ "device_summary_no_close", IOTDB_LOCAL_PKG + ".DeviceSummaryNoCloseFunction");
+ tableResultSetEqualTest(
+ "SELECT device_summary_no_close(device_id) AS summary FROM readings WHERE device_id = 'd1'",
+ new String[] {"summary"},
+ new String[] {"一号车间温度传感器(上限:30.0),"},
+ DATABASE_NAME);
+ tableResultSetEqualTest(
+ "SELECT device_summary_no_close(device_id) AS summary FROM readings WHERE device_id = 'd2'",
+ new String[] {"summary"},
+ new String[] {"二号车间温度传感器(上限:35.0),"},
+ DATABASE_NAME);
+ }
+
+ @Test
+ public void testRejectMutatingStatementsViaIoTDBLocal() {
+ SQLFunctionUtils.createUDF("insert_probe", IOTDB_LOCAL_PKG + ".InsertProbeFunction");
+ tableAssertTestFail(
+ "SELECT insert_probe(device_id) FROM readings WHERE device_id = 'd1'",
+ "701: Only query is supported for IoTDBLocal query interface",
+ DATABASE_NAME);
+ tableResultSetEqualTest(
+ "SELECT COUNT(*) AS cnt FROM probe_table",
+ new String[] {"cnt"},
+ new String[] {"1,"},
+ DATABASE_NAME);
+
+ SQLFunctionUtils.createUDF("create_table_probe", IOTDB_LOCAL_PKG + ".CreateTableProbeFunction");
+ tableAssertTestFail(
+ "SELECT create_table_probe(device_id) FROM readings WHERE device_id = 'd1'",
+ "701: Only query is supported for IoTDBLocal query interface",
+ DATABASE_NAME);
+ tableAssertTestFail(
+ "SELECT * FROM should_not_exist_local_probe", "does not exist", DATABASE_NAME);
+
+ SQLFunctionUtils.createUDF("drop_table_probe", IOTDB_LOCAL_PKG + ".DropTableProbeFunction");
+ tableAssertTestFail(
+ "SELECT drop_table_probe(device_id) FROM readings WHERE device_id = 'd1'",
+ "701: Only query is supported for IoTDBLocal query interface",
+ DATABASE_NAME);
+ tableResultSetEqualTest(
+ "SELECT device_id FROM device_info ORDER BY device_id",
+ new String[] {"device_id"},
+ new String[] {"d1,", "d2,"},
+ DATABASE_NAME);
+ }
+
+ @Test
+ public void testLocalQueryUdafAtBeforeStart() {
+ SQLFunctionUtils.createUDF(
+ "local_query_udaf_before_start",
+ IOTDB_LOCAL_PKG + ".LocalQueryUdafAtBeforeStartAggregateFunction");
+ tableResultSetEqualTest(
+ "SELECT local_query_udaf_before_start(temperature) AS total FROM readings",
+ new String[] {"total"},
+ new String[] {"5,"},
+ DATABASE_NAME);
+ }
+
+ @Test
+ public void testLocalQueryGroupedUdafAtBeforeStart() {
+ SQLFunctionUtils.createUDF(
+ "local_query_grouped_udaf_before_start",
+ IOTDB_LOCAL_PKG + ".LocalQueryUdafAtBeforeStartAggregateFunction");
+ // GROUP BY device_id, temperature disables agg pushdown and exercises
+ // GroupedUserDefinedAggregateAccumulator (hash aggregation path).
+ tableResultSetEqualTest(
+ "SELECT device_id, local_query_grouped_udaf_before_start(temperature) AS total "
+ + "FROM readings GROUP BY device_id, temperature ORDER BY device_id",
+ new String[] {"device_id", "total"},
+ new String[] {"d1,3,", "d2,3,", "d3,3,"},
+ DATABASE_NAME);
+ }
+
+ @Test
+ public void testLocalQueryGroupedUdafAtBeforeStartViaAggPushdown() {
+ SQLFunctionUtils.createUDF(
+ "local_query_grouped_udaf_agg_pushdown",
+ IOTDB_LOCAL_PKG + ".LocalQueryUdafAtBeforeStartAggregateFunction");
+ // GROUP BY tag only: aggregation is pushed into AggTableScan, one device per group.
+ tableResultSetEqualTest(
+ "SELECT device_id, local_query_grouped_udaf_agg_pushdown(temperature) AS total "
+ + "FROM readings GROUP BY device_id ORDER BY device_id",
+ new String[] {"device_id", "total"},
+ new String[] {"d1,3,", "d2,3,", "d3,3,"},
+ DATABASE_NAME);
+ }
+
+ @Test
+ public void testLocalQueryUdafInAddInput() {
+ SQLFunctionUtils.createUDF(
+ "local_query_udaf_add_input",
+ IOTDB_LOCAL_PKG + ".LocalQueryUdafInAddInputAggregateFunction");
+ tableResultSetEqualTest(
+ "SELECT local_query_udaf_add_input(temperature) AS total FROM readings",
+ new String[] {"total"},
+ new String[] {"5,"},
+ DATABASE_NAME);
+ }
+
+ @Test
+ public void testLocalQueryUdafInCombineState() {
+ assumeTrue(
+ "combineState IoTDBLocal.query coverage requires cluster execution",
+ EnvFactory.getEnv().getDataNodeWrapperList().size() > 1);
+ SQLFunctionUtils.createUDF(
+ "local_query_udaf_combine_state",
+ IOTDB_LOCAL_PKG + ".LocalQueryUdafInCombineStateAggregateFunction");
+ tableResultSetEqualTest(
+ "SELECT local_query_udaf_combine_state(temperature) AS total FROM readings",
+ new String[] {"total"},
+ new String[] {"5,"},
+ DATABASE_NAME);
+ }
+
+ @Test
+ public void testLocalQueryUdafInOutputFinal() {
+ SQLFunctionUtils.createUDF(
+ "local_query_udaf_output_final",
+ IOTDB_LOCAL_PKG + ".LocalQueryUdafInOutputFinalAggregateFunction");
+ tableResultSetEqualTest(
+ "SELECT local_query_udaf_output_final(temperature) AS total FROM readings",
+ new String[] {"total"},
+ new String[] {"5,"},
+ DATABASE_NAME);
+ }
+
+ @Test
+ public void testLocalQueryUdafAtBeforeDestroy() {
+ SQLFunctionUtils.createUDF(
+ "local_query_udaf_before_destroy",
+ IOTDB_LOCAL_PKG + ".LocalQueryUdafAtBeforeDestroyAggregateFunction");
+ tableResultSetEqualTest(
+ "SELECT local_query_udaf_before_destroy(temperature) AS total FROM readings",
+ new String[] {"total"},
+ new String[] {"3,"},
+ DATABASE_NAME);
+ }
+
+ @Test
+ public void testLocalQueryLeafTableFunctionAtBeforeStart() {
+ SQLFunctionUtils.createUDF("device_id_list", IOTDB_LOCAL_PKG + ".DeviceIdListTableFunction");
+ tableResultSetEqualTest(
+ "SELECT * FROM device_id_list('id:')",
+ new String[] {"device_id"},
+ new String[] {"id:d1,", "id:d2,"},
+ DATABASE_NAME);
+ }
+
+ @Test
+ public void testLocalQueryLeafTableFunctionInProcess() {
+ SQLFunctionUtils.createUDF(
+ "device_id_list_in_process", IOTDB_LOCAL_PKG + ".DeviceIdListInProcessTableFunction");
+ tableResultSetEqualTest(
+ "SELECT * FROM device_id_list_in_process('id:')",
+ new String[] {"device_id"},
+ new String[] {"id:d1,", "id:d2,"},
+ DATABASE_NAME);
+ }
+
+ @Test
+ public void testLocalQueryDataTableFunctionAtBeforeStart() {
+ SQLFunctionUtils.createUDF(
+ "enrich_device_name_at_before_start",
+ IOTDB_LOCAL_PKG + ".DeviceNameEnrichBeforeStartTableFunction");
+ tableResultSetEqualTest(
+ "SELECT device_name FROM enrich_device_name_at_before_start((SELECT device_id FROM readings ORDER BY time)) order by device_name",
+ new String[] {"device_name"},
+ new String[] {
+ "一号车间温度传感器,", "二号车间温度传感器,", "未知设备,",
+ },
+ DATABASE_NAME);
+ }
+
+ @Test
+ public void testLocalQueryDataTableFunctionInProcess() {
+ SQLFunctionUtils.createUDF(
+ "enrich_device_name_in_process",
+ IOTDB_LOCAL_PKG + ".DeviceNameEnrichInProcessTableFunction");
+ tableResultSetEqualTest(
+ "SELECT device_name FROM enrich_device_name_in_process((SELECT device_id FROM readings ORDER BY time)) order by device_name",
+ new String[] {"device_name"},
+ new String[] {
+ "一号车间温度传感器,", "二号车间温度传感器,", "未知设备,",
+ },
+ DATABASE_NAME);
+ }
+
+ // ── permission inheritance ──────────────────────────────────────────────────
+
+ @Test
+ public void testIoTDBLocalInheritsSelectPermission() {
+ setupLimitedUserWithTableGrants("readings", "device_info", "device_limits");
+ SQLFunctionUtils.createUDF("device_name", IOTDB_LOCAL_PKG + ".DeviceNameFunction");
+ tableResultSetEqualTest(
+ "SELECT device_name(device_id) AS name FROM readings WHERE device_id = 'd1'",
+ new String[] {"name"},
+ new String[] {"一号车间温度传感器,"},
+ LIMITED_USER,
+ LIMITED_PASSWORD,
+ DATABASE_NAME);
+ dropLimitedUser();
+ }
+
+ @Test
+ public void testIoTDBLocalDeniedWithoutTablePermission() {
+ setupLimitedUserWithTableGrants("readings");
+ SQLFunctionUtils.createUDF("secret_query", IOTDB_LOCAL_PKG + ".SecretTableQueryFunction");
+ tableAssertTestFail(
+ "SELECT secret_query(device_id) FROM readings WHERE device_id = 'd1'",
+ "Access Denied",
+ LIMITED_USER,
+ LIMITED_PASSWORD,
+ DATABASE_NAME);
+ dropLimitedUser();
+ }
+
+ private static void setupLimitedUserWithTableGrants(String... tables) {
+ try (Connection connection = EnvFactory.getEnv().getTableConnection();
+ Statement statement = connection.createStatement()) {
+ try {
+ statement.execute("DROP USER " + LIMITED_USER);
+ } catch (Exception ignored) {
+ // user may not exist
+ }
+ statement.execute("CREATE USER " + LIMITED_USER + " '" + LIMITED_PASSWORD + "'");
+ for (String table : tables) {
+ statement.execute(
+ "GRANT SELECT ON " + DATABASE_NAME + "." + table + " TO USER " + LIMITED_USER);
+ }
+ } catch (Exception e) {
+ fail("setupLimitedUserWithTableGrants failed: " + e.getMessage());
+ }
+ }
+
+ private static void dropLimitedUser() {
+ try (Connection connection = EnvFactory.getEnv().getTableConnection();
+ Statement statement = connection.createStatement()) {
+ try {
+ statement.execute("DROP USER " + LIMITED_USER);
+ } catch (Exception ignored) {
+ // user may not exist
+ }
+ } catch (Exception e) {
+ fail("dropLimitedUser failed: " + e.getMessage());
+ }
+ }
+}
diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBLocalLogIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBLocalLogIT.java
new file mode 100644
index 00000000000..c045357296a
--- /dev/null
+++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBLocalLogIT.java
@@ -0,0 +1,159 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.relational.it.db.it.udf;
+
+import org.apache.iotdb.db.query.udf.example.relational.iotdblocal.IoTDBLocalLogHelper;
+import org.apache.iotdb.it.env.EnvFactory;
+import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper;
+import org.apache.iotdb.it.framework.IoTDBTestRunner;
+import org.apache.iotdb.itbase.category.TableLocalStandaloneIT;
+
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Assert;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+
+import java.io.IOException;
+import java.sql.Connection;
+import java.sql.Statement;
+
+import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest;
+import static org.junit.Assert.fail;
+
+/**
+ * Verifies {@link org.apache.iotdb.udf.api.IoTDBLocal} logging APIs on DataNode log output. Runs
+ * only in 1C1D ({@link TableLocalStandaloneIT}) because log inspection uses {@code
+ * DataNodeWrapper(0)}.
+ */
+@RunWith(IoTDBTestRunner.class)
+@Category({TableLocalStandaloneIT.class})
+public class IoTDBLocalLogIT {
+
+ private static final String DATABASE_NAME = "iotdb_local_log_it";
+ private static final String IOTDB_LOCAL_PKG =
+ "org.apache.iotdb.db.query.udf.example.relational.iotdblocal";
+
+ private static DataNodeWrapper dataNodeWrapper;
+
+ private static final String[] SETUP_SQLS =
+ new String[] {
+ "CREATE DATABASE " + DATABASE_NAME,
+ "USE " + DATABASE_NAME,
+ "CREATE TABLE readings (device_id STRING TAG, temperature DOUBLE FIELD)",
+ "CREATE TABLE vehicle (device_id STRING TAG, s1 INT32 FIELD)",
+ "INSERT INTO readings(time, device_id, temperature) VALUES (1000, 'd1', 25.5)",
+ "INSERT INTO vehicle(time, device_id, s1) VALUES (1, 'd0', 1), (2, 'd0', 2), (3, 'd0', 3)",
+ "FLUSH",
+ };
+
+ @BeforeClass
+ public static void setUp() throws Exception {
+ EnvFactory.getEnv().initClusterEnvironment();
+ dataNodeWrapper = EnvFactory.getEnv().getDataNodeWrapper(0);
+ executeAsRoot(SETUP_SQLS);
+ }
+
+ @AfterClass
+ public static void tearDown() throws Exception {
+ EnvFactory.getEnv().cleanClusterEnvironment();
+ }
+
+ @After
+ public void dropFunctions() {
+ SQLFunctionUtils.dropAllUDF();
+ }
+
+ @Test
+ public void testScalarFunctionLogApis() throws IOException {
+ SQLFunctionUtils.createUDF(
+ "iotdb_local_log_scalar", IOTDB_LOCAL_PKG + ".IoTDBLocalLogFunction");
+ dataNodeWrapper.clearLogContent();
+ tableResultSetEqualTest(
+ "SELECT iotdb_local_log_scalar(device_id) AS log_ok FROM readings WHERE device_id = 'd1'",
+ new String[] {"log_ok"},
+ new String[] {"ok,"},
+ DATABASE_NAME);
+ assertPhaseLogs(IoTDBLocalLogHelper.SCALAR_BEFORE_START);
+ assertPhaseLogs(IoTDBLocalLogHelper.SCALAR_EVALUATE);
+ assertPhaseLogs(IoTDBLocalLogHelper.SCALAR_BEFORE_DESTROY);
+ }
+
+ @Test
+ public void testAggregateFunctionLogApis() throws IOException {
+ SQLFunctionUtils.createUDF(
+ "iotdb_local_log_udaf", IOTDB_LOCAL_PKG + ".IoTDBLocalLogAggregateFunction");
+ dataNodeWrapper.clearLogContent();
+ tableResultSetEqualTest(
+ "SELECT iotdb_local_log_udaf(s1) AS row_count FROM vehicle",
+ new String[] {"row_count"},
+ new String[] {"3,"},
+ DATABASE_NAME);
+ assertPhaseLogs(IoTDBLocalLogHelper.UDAF_BEFORE_START);
+ assertPhaseLogs(IoTDBLocalLogHelper.UDAF_ADD_INPUT);
+ assertPhaseLogs(IoTDBLocalLogHelper.UDAF_OUTPUT_FINAL);
+ assertPhaseLogs(IoTDBLocalLogHelper.UDAF_BEFORE_DESTROY);
+ }
+
+ @Test
+ public void testTableFunctionLogApis() throws IOException {
+ SQLFunctionUtils.createUDF(
+ "iotdb_local_log_split", IOTDB_LOCAL_PKG + ".IoTDBLocalLogTableFunction");
+ dataNodeWrapper.clearLogContent();
+ tableResultSetEqualTest(
+ "SELECT * FROM iotdb_local_log_split('a,b')",
+ new String[] {"output"},
+ new String[] {"a,", "b,"},
+ DATABASE_NAME);
+ assertPhaseLogs(IoTDBLocalLogHelper.TVF_BEFORE_START);
+ assertPhaseLogs(IoTDBLocalLogHelper.TVF_PROCESS);
+ assertPhaseLogs(IoTDBLocalLogHelper.TVF_BEFORE_DESTROY);
+ }
+
+ private static void assertPhaseLogs(String phase) throws IOException {
+ assertLogContains(IoTDBLocalLogHelper.infoPlain(phase));
+ assertLogContains(IoTDBLocalLogHelper.infoFormat(phase));
+ assertLogContains(IoTDBLocalLogHelper.infoCause(phase));
+ assertLogContains(IoTDBLocalLogHelper.warnPlain(phase));
+ assertLogContains(IoTDBLocalLogHelper.warnFormat(phase));
+ assertLogContains(IoTDBLocalLogHelper.warnCause(phase));
+ assertLogContains(IoTDBLocalLogHelper.errorPlain(phase));
+ assertLogContains(IoTDBLocalLogHelper.errorFormat(phase));
+ assertLogContains(IoTDBLocalLogHelper.errorCause(phase));
+ assertLogContains(IoTDBLocalLogHelper.CAUSE_MESSAGE);
+ }
+
+ private static void assertLogContains(String content) throws IOException {
+ Assert.assertTrue("Expected log to contain: " + content, dataNodeWrapper.logContains(content));
+ }
+
+ private static void executeAsRoot(String... sqls) {
+ try (Connection connection = EnvFactory.getEnv().getTableConnection();
+ Statement statement = connection.createStatement()) {
+ for (String sql : sqls) {
+ statement.execute(sql);
+ }
+ } catch (Exception e) {
+ fail("executeAsRoot failed: " + e.getMessage());
+ }
+ }
+}
diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/SQLFunctionUtils.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/SQLFunctionUtils.java
index 34f57b2940b..6174f06bb5b 100644
--- a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/SQLFunctionUtils.java
+++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/SQLFunctionUtils.java
@@ -41,7 +41,6 @@ public class SQLFunctionUtils {
public static void createUDF(String udfName, String classPath) {
try (Connection connection = EnvFactory.getEnv().getTableConnection();
Statement statement = connection.createStatement()) {
- // create
statement.execute(String.format("create function %s as '%s'", udfName, classPath));
// check
try (ResultSet resultSet = statement.executeQuery("show functions")) {
diff --git a/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/IoTDBLocal.java b/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/IoTDBLocal.java
new file mode 100644
index 00000000000..b84cb45965a
--- /dev/null
+++ b/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/IoTDBLocal.java
@@ -0,0 +1,70 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.udf.api;
+
+import org.apache.iotdb.udf.api.exception.UDFException;
+
+/**
+ * Entry point for UDF-embedded SQL queries and logging, injected by the framework during UDF
+ * lifecycle. Reuses the outer query caller's user identity and permissions without extra client
+ * dependencies.
+ */
+public interface IoTDBLocal {
+
+ /**
+ * Execute a single read-only table-model SQL statement and return a streaming result set.
+ *
+ * @param sql table-model read SQL (e.g. SELECT, SHOW CONFIGNODES)
+ * @return query result set; call {@link UDFResultSet#close()} when done or rely on framework
+ * cleanup
+ * @throws UDFException on parse failure, permission denied, or execution error
+ */
+ UDFResultSet query(String sql) throws UDFException;
+
+ /** Log at INFO level. */
+ void info(String msg);
+
+ /** Log at INFO level with formatting. */
+ void info(String format, Object... args);
+
+ /** Log at INFO level with exception stack. */
+ void info(String msg, Throwable t);
+
+ /** Log at WARN level. */
+ void warn(String msg);
+
+ /** Log at WARN level with formatting. */
+ void warn(String format, Object... args);
+
+ /** Log at WARN level with exception stack. */
+ void warn(String msg, Throwable t);
+
+ /** Log at ERROR level. */
+ void error(String msg);
+
+ /** Log at ERROR level with formatting. */
+ void error(String format, Object... args);
+
+ /** Log at ERROR level with exception stack. */
+ void error(String msg, Throwable t);
+
+ /** Close internal session. Called by the framework after beforeDestroy method. */
+ void close();
+}
diff --git a/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionLeafProcessor.java b/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/UDFResultSet.java
similarity index 51%
copy from iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionLeafProcessor.java
copy to iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/UDFResultSet.java
index ab8763f3c4a..ed6d13442bf 100644
--- a/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionLeafProcessor.java
+++ b/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/UDFResultSet.java
@@ -17,31 +17,29 @@
* under the License.
*/
-package org.apache.iotdb.udf.api.relational.table.processor;
+package org.apache.iotdb.udf.api;
-import org.apache.tsfile.block.column.ColumnBuilder;
+import org.apache.iotdb.udf.api.exception.UDFException;
+import org.apache.iotdb.udf.api.relational.access.Record;
-import java.util.List;
-
-public interface TableFunctionLeafProcessor {
-
- default void beforeStart() {
- // do nothing
- }
+/** Streaming result set returned by {@link IoTDBLocal#query(String)}. */
+public interface UDFResultSet extends AutoCloseable {
/**
- * This method processes a portion of data. It is called multiple times until the processor is
- * fully processed.
- *
- * @param columnBuilders a list of {@link ColumnBuilder} for each column in the output table.
+ * @return {@code true} if another row is available
+ * @throws UDFException if underlying read fails
*/
- void process(List<ColumnBuilder> columnBuilders);
+ boolean hasNext() throws UDFException;
- /** This method is called to determine if the processor has finished processing all data. */
- boolean isFinish();
+ /**
+ * @return the next row
+ * @throws UDFException if underlying read fails
+ * @throws java.util.NoSuchElementException if no more rows (consistent with {@link
+ * java.util.Iterator})
+ */
+ Record next() throws UDFException;
- /** This method is mainly used to release the resources used in the UDF. */
- default void beforeDestroy() {
- // do nothing
- }
+ /** Release query resources. Repeated calls must be idempotent. */
+ @Override
+ void close() throws UDFException;
}
diff --git a/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/AggregateFunction.java b/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/AggregateFunction.java
index 6c9d385ac24..6c29de6e7ad 100644
--- a/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/AggregateFunction.java
+++ b/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/AggregateFunction.java
@@ -19,6 +19,7 @@
package org.apache.iotdb.udf.api.relational;
+import org.apache.iotdb.udf.api.IoTDBLocal;
import org.apache.iotdb.udf.api.State;
import org.apache.iotdb.udf.api.customizer.analysis.AggregateFunctionAnalysis;
import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;
@@ -58,6 +59,26 @@ public interface AggregateFunction extends SQLFunction {
// do nothing
}
+ default void beforeStart(FunctionArguments arguments, IoTDBLocal local) throws UDFException {
+ beforeStart(arguments);
+ }
+
+ /**
+ * Same as {@link #addInput(State, Record)} with access to {@link IoTDBLocal} for embedded
+ * queries.
+ */
+ default void addInput(State state, Record input, IoTDBLocal local) {
+ addInput(state, input);
+ }
+
+ /**
+ * Same as {@link #combineState(State, State)} with access to {@link IoTDBLocal} for embedded
+ * queries.
+ */
+ default void combineState(State state, State rhs, IoTDBLocal local) {
+ combineState(state, rhs);
+ }
+
/** Create and initialize state. You may bind some resource in this method. */
State createState();
@@ -85,6 +106,14 @@ public interface AggregateFunction extends SQLFunction {
*/
void outputFinal(State state, ResultValue resultValue);
+ /**
+ * Same as {@link #outputFinal(State, ResultValue)} with access to {@link IoTDBLocal} for embedded
+ * queries.
+ */
+ default void outputFinal(State state, ResultValue resultValue, IoTDBLocal local) {
+ outputFinal(state, resultValue);
+ }
+
/**
* Remove input data from state. This method is used to remove the data points that have been
* added to the state. Once it is implemented, {@linkplain
@@ -101,4 +130,8 @@ public interface AggregateFunction extends SQLFunction {
default void beforeDestroy() {
// do nothing
}
+
+ default void beforeDestroy(IoTDBLocal local) {
+ beforeDestroy();
+ }
}
diff --git a/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/ScalarFunction.java b/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/ScalarFunction.java
index 68d7793093d..e1daab6e4b0 100644
--- a/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/ScalarFunction.java
+++ b/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/ScalarFunction.java
@@ -19,6 +19,7 @@
package org.apache.iotdb.udf.api.relational;
+import org.apache.iotdb.udf.api.IoTDBLocal;
import org.apache.iotdb.udf.api.customizer.analysis.ScalarFunctionAnalysis;
import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;
import org.apache.iotdb.udf.api.exception.UDFArgumentNotValidException;
@@ -54,6 +55,14 @@ public interface ScalarFunction extends SQLFunction {
// do nothing
}
+ /**
+ * Same as {@link #beforeStart(FunctionArguments)} with access to {@link IoTDBLocal} for embedded
+ * queries.
+ */
+ default void beforeStart(FunctionArguments arguments, IoTDBLocal local) throws UDFException {
+ beforeStart(arguments);
+ }
+
/**
* This method will be called to process the transformation. In a single UDF query, this method
* may be called multiple times.
@@ -63,8 +72,18 @@ public interface ScalarFunction extends SQLFunction {
*/
Object evaluate(Record input) throws UDFException;
+ /** Same as {@link #evaluate(Record)} with access to {@link IoTDBLocal} for embedded queries. */
+ default Object evaluate(Record input, IoTDBLocal local) throws UDFException {
+ return evaluate(input);
+ }
+
/** This method is mainly used to release the resources used in the ScalarFunction. */
default void beforeDestroy() {
// do nothing
}
+
+ /** Same as {@link #beforeDestroy()} with access to {@link IoTDBLocal}. */
+ default void beforeDestroy(IoTDBLocal local) {
+ beforeDestroy();
+ }
}
diff --git a/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionDataProcessor.java b/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionDataProcessor.java
index 0fe40fc894b..e4960fce409 100644
--- a/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionDataProcessor.java
+++ b/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionDataProcessor.java
@@ -19,6 +19,7 @@
package org.apache.iotdb.udf.api.relational.table.processor;
+import org.apache.iotdb.udf.api.IoTDBLocal;
import org.apache.iotdb.udf.api.relational.access.Record;
import org.apache.iotdb.udf.api.relational.table.TableFunctionAnalysis;
@@ -33,6 +34,11 @@ public interface TableFunctionDataProcessor {
// do nothing
}
+ /** Same as {@link #beforeStart()} with access to {@link IoTDBLocal} for embedded queries. */
+ default void beforeStart(IoTDBLocal local) {
+ beforeStart();
+ }
+
/**
* This method processes a portion of data. It is called multiple times until the partition is
* fully processed.
@@ -51,6 +57,18 @@ public interface TableFunctionDataProcessor {
List<ColumnBuilder> properColumnBuilders,
ColumnBuilder passThroughIndexBuilder);
+ /**
+ * Same as {@link #process(Record, List, ColumnBuilder)} with access to {@link IoTDBLocal} for
+ * embedded queries.
+ */
+ default void process(
+ Record input,
+ List<ColumnBuilder> properColumnBuilders,
+ ColumnBuilder passThroughIndexBuilder,
+ IoTDBLocal local) {
+ process(input, properColumnBuilders, passThroughIndexBuilder);
+ }
+
/**
* This method is called after all data is processed. It is used to finalize the output table and
* close resource. All remaining data should be written to the columnBuilders.
@@ -66,8 +84,24 @@ public interface TableFunctionDataProcessor {
// do nothing
}
+ /**
+ * Same as {@link #finish(List, ColumnBuilder)} with access to {@link IoTDBLocal} for embedded
+ * queries.
+ */
+ default void finish(
+ List<ColumnBuilder> properColumnBuilders,
+ ColumnBuilder passThroughIndexBuilder,
+ IoTDBLocal local) {
+ finish(properColumnBuilders, passThroughIndexBuilder);
+ }
+
/** This method is mainly used to release the resources used in the UDF. */
default void beforeDestroy() {
// do nothing
}
+
+ /** Same as {@link #beforeDestroy()} with access to {@link IoTDBLocal}. */
+ default void beforeDestroy(IoTDBLocal local) {
+ beforeDestroy();
+ }
}
diff --git a/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionLeafProcessor.java b/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionLeafProcessor.java
index ab8763f3c4a..6da0c8d5149 100644
--- a/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionLeafProcessor.java
+++ b/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/relational/table/processor/TableFunctionLeafProcessor.java
@@ -19,6 +19,8 @@
package org.apache.iotdb.udf.api.relational.table.processor;
+import org.apache.iotdb.udf.api.IoTDBLocal;
+
import org.apache.tsfile.block.column.ColumnBuilder;
import java.util.List;
@@ -29,6 +31,11 @@ public interface TableFunctionLeafProcessor {
// do nothing
}
+ /** Same as {@link #beforeStart()} with access to {@link IoTDBLocal} for embedded queries. */
+ default void beforeStart(IoTDBLocal local) {
+ beforeStart();
+ }
+
/**
* This method processes a portion of data. It is called multiple times until the processor is
* fully processed.
@@ -37,6 +44,11 @@ public interface TableFunctionLeafProcessor {
*/
void process(List<ColumnBuilder> columnBuilders);
+ /** Same as {@link #process(List)} with access to {@link IoTDBLocal} for embedded queries. */
+ default void process(List<ColumnBuilder> columnBuilders, IoTDBLocal local) {
+ process(columnBuilders);
+ }
+
/** This method is called to determine if the processor has finished processing all data. */
boolean isFinish();
@@ -44,4 +56,9 @@ public interface TableFunctionLeafProcessor {
default void beforeDestroy() {
// do nothing
}
+
+ /** Same as {@link #beforeDestroy()} with access to {@link IoTDBLocal}. */
+ default void beforeDestroy(IoTDBLocal local) {
+ beforeDestroy();
+ }
}
diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/function/TableFunctionLeafOperator.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/function/TableFunctionLeafOperator.java
index ced20f7bbf0..589e175e146 100644
--- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/function/TableFunctionLeafOperator.java
+++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/function/TableFunctionLeafOperator.java
@@ -25,6 +25,7 @@ import org.apache.iotdb.calc.i18n.CalcMessages;
import org.apache.iotdb.calc.plan.planner.CommonOperatorUtils;
import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
import org.apache.iotdb.rpc.TSStatusCode;
+import org.apache.iotdb.udf.api.IoTDBLocal;
import org.apache.iotdb.udf.api.relational.table.TableFunctionProcessorProvider;
import org.apache.iotdb.udf.api.relational.table.processor.TableFunctionLeafProcessor;
@@ -40,6 +41,8 @@ import org.slf4j.LoggerFactory;
import java.util.Arrays;
import java.util.List;
+import static com.google.common.base.Preconditions.checkArgument;
+
// only one input source is supported now
public class TableFunctionLeafOperator implements ProcessOperator {
private static final Logger LOGGER = LoggerFactory.getLogger(TableFunctionLeafOperator.class);
@@ -48,22 +51,26 @@ public class TableFunctionLeafOperator implements ProcessOperator {
private final TsBlockBuilder blockBuilder;
private final TableFunctionLeafProcessor processor;
+ private final IoTDBLocal ioTDBLocal;
private volatile boolean init = false;
public TableFunctionLeafOperator(
CommonOperatorContext operatorContext,
TableFunctionProcessorProvider processorProvider,
- List<TSDataType> outputDataTypes) {
+ List<TSDataType> outputDataTypes,
+ IoTDBLocal ioTDBLocal) {
+ checkArgument(ioTDBLocal != null, "IoTDBLocal must not be null for table function");
this.operatorContext = operatorContext;
this.processor = processorProvider.getSplitProcessor();
this.blockBuilder = new TsBlockBuilder(outputDataTypes);
+ this.ioTDBLocal = ioTDBLocal;
}
@Override
public ListenableFuture<?> isBlocked() {
if (!init) {
init = true;
- processor.beforeStart();
+ processor.beforeStart(ioTDBLocal);
}
return NOT_BLOCKED;
}
@@ -77,7 +84,7 @@ public class TableFunctionLeafOperator implements ProcessOperator {
public TsBlock next() throws Exception {
List<ColumnBuilder> columnBuilders = getOutputColumnBuilders();
try {
- processor.process(columnBuilders);
+ processor.process(columnBuilders, ioTDBLocal);
} catch (Exception e) {
LOGGER.warn(CalcMessages.EXCEPTION_HAPPENED_WHEN_EXECUTING_UDTF, e);
throw new IoTDBRuntimeException(
@@ -106,7 +113,8 @@ public class TableFunctionLeafOperator implements ProcessOperator {
@Override
public void close() throws Exception {
try {
- processor.beforeDestroy();
+ processor.beforeDestroy(ioTDBLocal);
+ ioTDBLocal.close();
} catch (Exception e) {
LOGGER.warn(CalcMessages.EXCEPTION_HAPPENED_WHEN_EXECUTING_UDTF, e);
throw new IoTDBRuntimeException(
diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/function/TableFunctionOperator.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/function/TableFunctionOperator.java
index e4f11d7cd49..d4bc8a0da4c 100644
--- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/function/TableFunctionOperator.java
+++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/function/TableFunctionOperator.java
@@ -28,6 +28,7 @@ import org.apache.iotdb.calc.execution.operator.process.function.partition.Parti
import org.apache.iotdb.calc.execution.operator.process.function.partition.Slice;
import org.apache.iotdb.calc.plan.planner.CommonOperatorUtils;
import org.apache.iotdb.commons.queryengine.execution.MemoryEstimationHelper;
+import org.apache.iotdb.udf.api.IoTDBLocal;
import org.apache.iotdb.udf.api.relational.access.Record;
import org.apache.iotdb.udf.api.relational.table.TableFunctionProcessorProvider;
import org.apache.iotdb.udf.api.relational.table.processor.TableFunctionDataProcessor;
@@ -53,6 +54,8 @@ import java.util.List;
import java.util.Optional;
import java.util.Queue;
+import static com.google.common.base.Preconditions.checkArgument;
+
// only one input source is supported now
public class TableFunctionOperator implements ProcessOperator {
@@ -72,6 +75,7 @@ public class TableFunctionOperator implements ProcessOperator {
private final PartitionCache partitionCache;
private final boolean requireRecordSnapshot;
private final boolean isDeclaredAsPassThrough;
+ private final IoTDBLocal ioTDBLocal;
private TableFunctionDataProcessor processor;
private PartitionState partitionState;
@@ -91,7 +95,9 @@ public class TableFunctionOperator implements ProcessOperator {
List<Integer> passThroughChannels,
boolean isDeclaredAsPassThrough,
List<Integer> partitionChannels,
- boolean requireRecordSnapshot) {
+ boolean requireRecordSnapshot,
+ IoTDBLocal ioTDBLocal) {
+ checkArgument(ioTDBLocal != null, "IoTDBLocal must not be null for table function");
this.operatorContext = operatorContext;
this.inputOperator = inputOperator;
this.properChannelCount = properChannelCount;
@@ -106,6 +112,7 @@ public class TableFunctionOperator implements ProcessOperator {
this.partitionCache = new PartitionCache();
this.resultTsBlocks = new LinkedList<>();
this.requireRecordSnapshot = requireRecordSnapshot;
+ this.ioTDBLocal = ioTDBLocal;
}
@Override
@@ -159,7 +166,7 @@ public class TableFunctionOperator implements ProcessOperator {
ColumnBuilder passThroughIndexBuilder = getPassThroughIndexBuilder();
if (stateType == PartitionState.StateType.FINISHED) {
if (processor != null) {
- processor.finish(properColumnBuilders, passThroughIndexBuilder);
+ processor.finish(properColumnBuilders, passThroughIndexBuilder, ioTDBLocal);
}
finished = true;
resultTsBlocks.addAll(buildTsBlock(properColumnBuilders, passThroughIndexBuilder));
@@ -170,21 +177,22 @@ public class TableFunctionOperator implements ProcessOperator {
if (stateType == PartitionState.StateType.NEW_PARTITION) {
if (processor != null) {
// previous partition state has not finished consuming yet
- processor.finish(properColumnBuilders, passThroughIndexBuilder);
+ processor.finish(properColumnBuilders, passThroughIndexBuilder, ioTDBLocal);
resultTsBlocks.addAll(buildTsBlock(properColumnBuilders, passThroughIndexBuilder));
partitionCache.clear();
- processor.beforeDestroy();
+ destroyProcessor(processor);
processor = null;
return resultTsBlocks.poll();
} else {
processor = processorProvider.getDataProcessor();
- processor.beforeStart();
+ processor.beforeStart(ioTDBLocal);
}
}
partitionCache.addSlice(slice);
Iterator<Record> recordIterator = slice.getRequiredRecordIterator(requireRecordSnapshot);
while (recordIterator.hasNext()) {
- processor.process(recordIterator.next(), properColumnBuilders, passThroughIndexBuilder);
+ processor.process(
+ recordIterator.next(), properColumnBuilders, passThroughIndexBuilder, ioTDBLocal);
}
consumeCurrentPartitionState();
resultTsBlocks.addAll(buildTsBlock(properColumnBuilders, passThroughIndexBuilder));
@@ -251,6 +259,10 @@ public class TableFunctionOperator implements ProcessOperator {
isBlocked = null;
}
+ private void destroyProcessor(TableFunctionDataProcessor dataProcessor) {
+ dataProcessor.beforeDestroy(ioTDBLocal);
+ }
+
@Override
public boolean hasNext() throws Exception {
return !finished || !resultTsBlocks.isEmpty();
@@ -261,8 +273,10 @@ public class TableFunctionOperator implements ProcessOperator {
partitionCache.close();
inputOperator.close();
if (processor != null) {
- processor.beforeDestroy();
+ destroyProcessor(processor);
+ processor = null;
}
+ ioTDBLocal.close();
}
@Override
diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/AccumulatorFactory.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/AccumulatorFactory.java
index a37f1bef0c6..f9d35cf67f5 100644
--- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/AccumulatorFactory.java
+++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/AccumulatorFactory.java
@@ -65,6 +65,7 @@ import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Expression;
import org.apache.iotdb.commons.queryengine.plan.relational.type.InternalTypeManager;
import org.apache.iotdb.commons.queryengine.plan.udf.TableUDFUtils;
import org.apache.iotdb.commons.udf.utils.UDFDataTypeTransformer;
+import org.apache.iotdb.udf.api.IoTDBLocal;
import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;
import org.apache.iotdb.udf.api.relational.AggregateFunction;
@@ -78,6 +79,8 @@ import org.apache.tsfile.read.common.type.Type;
import org.apache.tsfile.read.common.type.TypeFactory;
import org.apache.tsfile.write.UnSupportedDataTypeException;
+import javax.annotation.Nullable;
+
import java.util.Arrays;
import java.util.List;
import java.util.Map;
@@ -86,6 +89,7 @@ import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;
+import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
import static java.util.Objects.requireNonNull;
import static org.apache.iotdb.calc.plan.relational.planner.ir.GlobalTimePredicateExtractVisitor.isMeasurementColumn;
@@ -110,13 +114,41 @@ public class AccumulatorFactory {
Set<String> measurementColumnNames,
boolean distinct,
MemoryReservationManager memoryReservationManager) {
+ return createAccumulator(
+ functionName,
+ aggregationType,
+ inputDataTypes,
+ inputExpressions,
+ inputAttributes,
+ ascending,
+ isAggTableScan,
+ timeColumnName,
+ measurementColumnNames,
+ distinct,
+ memoryReservationManager,
+ null);
+ }
+
+ public static TableAccumulator createAccumulator(
+ String functionName,
+ TAggregationType aggregationType,
+ List<TSDataType> inputDataTypes,
+ List<Expression> inputExpressions,
+ Map<String, String> inputAttributes,
+ boolean ascending,
+ boolean isAggTableScan,
+ String timeColumnName,
+ Set<String> measurementColumnNames,
+ boolean distinct,
+ MemoryReservationManager memoryReservationManager,
+ @Nullable IoTDBLocal ioTDBLocal) {
TableAccumulator result;
// Input expression size of 1 indicates aggregation split has occurred and this is a final
// aggregation
if (aggregationType == TAggregationType.UDAF) {
// If UDAF accumulator receives raw input, it needs to check input's attribute
- result = createUDAFAccumulator(functionName, inputDataTypes, inputAttributes);
+ result = createUDAFAccumulator(functionName, inputDataTypes, inputAttributes, ioTDBLocal);
} else if ((LAST_BY.getFunctionName().equals(functionName)
|| FIRST_BY.getFunctionName().equals(functionName))
&& inputExpressions.size() > 1) {
@@ -193,11 +225,34 @@ public class AccumulatorFactory {
boolean ascending,
boolean distinct,
MemoryReservationManager memoryReservationManager) {
+ return createGroupedAccumulator(
+ functionName,
+ aggregationType,
+ inputDataTypes,
+ inputExpressions,
+ inputAttributes,
+ ascending,
+ distinct,
+ memoryReservationManager,
+ null);
+ }
+
+ public static GroupedAccumulator createGroupedAccumulator(
+ String functionName,
+ TAggregationType aggregationType,
+ List<TSDataType> inputDataTypes,
+ List<Expression> inputExpressions,
+ Map<String, String> inputAttributes,
+ boolean ascending,
+ boolean distinct,
+ MemoryReservationManager memoryReservationManager,
+ @Nullable IoTDBLocal ioTDBLocal) {
GroupedAccumulator result;
if (aggregationType == TAggregationType.UDAF) {
// If UDAF accumulator receives raw input, it needs to check input's attribute
- result = createGroupedUDAFAccumulator(functionName, inputDataTypes, inputAttributes);
+ result =
+ createGroupedUDAFAccumulator(functionName, inputDataTypes, inputAttributes, ioTDBLocal);
} else {
result =
createBuiltinGroupedAccumulator(
@@ -222,28 +277,38 @@ public class AccumulatorFactory {
}
private static TableAccumulator createUDAFAccumulator(
- String functionName, List<TSDataType> inputDataTypes, Map<String, String> inputAttributes) {
+ String functionName,
+ List<TSDataType> inputDataTypes,
+ Map<String, String> inputAttributes,
+ IoTDBLocal ioTDBLocal) {
+ checkArgument(ioTDBLocal != null, "IoTDBLocal must not be null for UDAF");
AggregateFunction aggregateFunction = TableUDFUtils.getAggregateFunction(functionName);
FunctionArguments functionArguments =
new FunctionArguments(
UDFDataTypeTransformer.transformToUDFDataTypeList(inputDataTypes), inputAttributes);
- aggregateFunction.beforeStart(functionArguments);
return new UserDefinedAggregateFunctionAccumulator(
aggregateFunction.analyze(functionArguments),
aggregateFunction,
- inputDataTypes.stream().map(TypeFactory::getType).collect(Collectors.toList()));
+ functionArguments,
+ inputDataTypes.stream().map(TypeFactory::getType).collect(Collectors.toList()),
+ ioTDBLocal);
}
private static GroupedAccumulator createGroupedUDAFAccumulator(
- String functionName, List<TSDataType> inputDataTypes, Map<String, String> inputAttributes) {
+ String functionName,
+ List<TSDataType> inputDataTypes,
+ Map<String, String> inputAttributes,
+ IoTDBLocal ioTDBLocal) {
+ checkArgument(ioTDBLocal != null, "IoTDBLocal must not be null for UDAF");
AggregateFunction aggregateFunction = TableUDFUtils.getAggregateFunction(functionName);
FunctionArguments functionArguments =
new FunctionArguments(
UDFDataTypeTransformer.transformToUDFDataTypeList(inputDataTypes), inputAttributes);
- aggregateFunction.beforeStart(functionArguments);
return new GroupedUserDefinedAggregateAccumulator(
aggregateFunction,
- inputDataTypes.stream().map(TypeFactory::getType).collect(Collectors.toList()));
+ functionArguments,
+ inputDataTypes.stream().map(TypeFactory::getType).collect(Collectors.toList()),
+ ioTDBLocal);
}
private static GroupedAccumulator createBuiltinGroupedAccumulator(
diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/UserDefinedAggregateFunctionAccumulator.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/UserDefinedAggregateFunctionAccumulator.java
index cd6d0dc1101..13c0b92305f 100644
--- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/UserDefinedAggregateFunctionAccumulator.java
+++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/UserDefinedAggregateFunctionAccumulator.java
@@ -20,8 +20,12 @@
package org.apache.iotdb.calc.execution.operator.source.relational.aggregation;
import org.apache.iotdb.calc.i18n.CalcMessages;
+import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
+import org.apache.iotdb.udf.api.IoTDBLocal;
import org.apache.iotdb.udf.api.State;
import org.apache.iotdb.udf.api.customizer.analysis.AggregateFunctionAnalysis;
+import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;
+import org.apache.iotdb.udf.api.exception.UDFException;
import org.apache.iotdb.udf.api.relational.AggregateFunction;
import org.apache.iotdb.udf.api.utils.ResultValue;
@@ -39,6 +43,7 @@ import java.util.Arrays;
import java.util.List;
import static com.google.common.base.Preconditions.checkArgument;
+import static org.apache.iotdb.rpc.TSStatusCode.EXECUTE_UDF_ERROR;
public class UserDefinedAggregateFunctionAccumulator implements TableAccumulator {
@@ -46,17 +51,50 @@ public class UserDefinedAggregateFunctionAccumulator implements TableAccumulator
RamUsageEstimator.shallowSizeOfInstance(UserDefinedAggregateFunctionAccumulator.class);
private final AggregateFunctionAnalysis analysis;
private final AggregateFunction aggregateFunction;
+ private final FunctionArguments functionArguments;
private final List<Type> inputDataTypes;
- private final State state;
+ private State state;
+ private final IoTDBLocal ioTDBLocal;
+ private boolean init;
public UserDefinedAggregateFunctionAccumulator(
AggregateFunctionAnalysis analysis,
AggregateFunction aggregateFunction,
- List<Type> inputDataTypes) {
+ FunctionArguments functionArguments,
+ List<Type> inputDataTypes,
+ IoTDBLocal ioTDBLocal) {
+ this(analysis, aggregateFunction, functionArguments, inputDataTypes, ioTDBLocal, false);
+ }
+
+ private UserDefinedAggregateFunctionAccumulator(
+ AggregateFunctionAnalysis analysis,
+ AggregateFunction aggregateFunction,
+ FunctionArguments functionArguments,
+ List<Type> inputDataTypes,
+ IoTDBLocal ioTDBLocal,
+ boolean init) {
+ checkArgument(ioTDBLocal != null, "IoTDBLocal must not be null for UDAF");
this.analysis = analysis;
this.aggregateFunction = aggregateFunction;
+ this.functionArguments = functionArguments;
this.inputDataTypes = inputDataTypes;
- this.state = aggregateFunction.createState();
+ this.ioTDBLocal = ioTDBLocal;
+ this.init = init;
+ this.state = init ? aggregateFunction.createState() : null;
+ }
+
+ private void initIfNeeded() {
+ if (init) {
+ return;
+ }
+ try {
+ aggregateFunction.beforeStart(functionArguments, ioTDBLocal);
+ init = true;
+ // create State after beforeStart
+ state = aggregateFunction.createState();
+ } catch (UDFException e) {
+ throw new IoTDBRuntimeException(e, EXECUTE_UDF_ERROR.getStatusCode());
+ }
}
@Override
@@ -66,23 +104,26 @@ public class UserDefinedAggregateFunctionAccumulator implements TableAccumulator
@Override
public TableAccumulator copy() {
- return new UserDefinedAggregateFunctionAccumulator(analysis, aggregateFunction, inputDataTypes);
+ return new UserDefinedAggregateFunctionAccumulator(
+ analysis, aggregateFunction, functionArguments, inputDataTypes, ioTDBLocal, init);
}
@Override
public void addInput(Column[] arguments, AggregationMask mask) {
+ initIfNeeded();
RecordIterator iterator =
mask.isSelectAll()
? new RecordIterator(
Arrays.asList(arguments), inputDataTypes, arguments[0].getPositionCount())
: new MaskedRecordIterator(Arrays.asList(arguments), inputDataTypes, mask);
while (iterator.hasNext()) {
- aggregateFunction.addInput(state, iterator.next());
+ aggregateFunction.addInput(state, iterator.next(), ioTDBLocal);
}
}
@Override
public void addIntermediate(Column argument) {
+ initIfNeeded();
checkArgument(
argument instanceof BinaryColumn
|| (argument instanceof RunLengthEncodedColumn
@@ -93,12 +134,13 @@ public class UserDefinedAggregateFunctionAccumulator implements TableAccumulator
otherState.reset();
Binary otherStateBinary = argument.getBinary(i);
otherState.deserialize(otherStateBinary.getValues());
- aggregateFunction.combineState(state, otherState);
+ aggregateFunction.combineState(state, otherState, ioTDBLocal);
}
}
@Override
public void evaluateIntermediate(ColumnBuilder columnBuilder) {
+ initIfNeeded();
checkArgument(
columnBuilder instanceof BinaryColumnBuilder,
"intermediate input and output of UDAF should be BinaryColumn");
@@ -108,8 +150,9 @@ public class UserDefinedAggregateFunctionAccumulator implements TableAccumulator
@Override
public void evaluateFinal(ColumnBuilder columnBuilder) {
+ initIfNeeded();
ResultValue resultValue = new ResultValue(columnBuilder);
- aggregateFunction.outputFinal(state, resultValue);
+ aggregateFunction.outputFinal(state, resultValue, ioTDBLocal);
}
@Override
@@ -126,6 +169,9 @@ public class UserDefinedAggregateFunctionAccumulator implements TableAccumulator
@Override
public void reset() {
+ if (!init) {
+ return;
+ }
state.reset();
}
@@ -150,7 +196,10 @@ public class UserDefinedAggregateFunctionAccumulator implements TableAccumulator
@Override
public void close() {
- aggregateFunction.beforeDestroy();
+ // ensure beforeStart was called
+ initIfNeeded();
+ aggregateFunction.beforeDestroy(ioTDBLocal);
+ ioTDBLocal.close();
state.destroyState();
}
}
diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/grouped/GroupedUserDefinedAggregateAccumulator.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/grouped/GroupedUserDefinedAggregateAccumulator.java
index 2ef228b0662..44df7fab1d7 100644
--- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/grouped/GroupedUserDefinedAggregateAccumulator.java
+++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/grouped/GroupedUserDefinedAggregateAccumulator.java
@@ -24,7 +24,11 @@ import org.apache.iotdb.calc.execution.operator.source.relational.aggregation.Ma
import org.apache.iotdb.calc.execution.operator.source.relational.aggregation.RecordIterator;
import org.apache.iotdb.calc.execution.operator.source.relational.aggregation.grouped.array.ObjectBigArray;
import org.apache.iotdb.calc.i18n.CalcMessages;
+import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
+import org.apache.iotdb.udf.api.IoTDBLocal;
import org.apache.iotdb.udf.api.State;
+import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;
+import org.apache.iotdb.udf.api.exception.UDFException;
import org.apache.iotdb.udf.api.relational.AggregateFunction;
import org.apache.iotdb.udf.api.utils.ResultValue;
@@ -41,20 +45,42 @@ import java.util.Arrays;
import java.util.List;
import static com.google.common.base.Preconditions.checkArgument;
+import static org.apache.iotdb.rpc.TSStatusCode.EXECUTE_UDF_ERROR;
public class GroupedUserDefinedAggregateAccumulator implements GroupedAccumulator {
private static final long INSTANCE_SIZE =
RamUsageEstimator.shallowSizeOfInstance(GroupedUserDefinedAggregateAccumulator.class);
private final AggregateFunction aggregateFunction;
+ private final FunctionArguments functionArguments;
private final ObjectBigArray<State> stateArray;
private final List<Type> inputDataTypes;
+ private final IoTDBLocal ioTDBLocal;
+ private boolean init = false;
public GroupedUserDefinedAggregateAccumulator(
- AggregateFunction aggregateFunction, List<Type> inputDataTypes) {
+ AggregateFunction aggregateFunction,
+ FunctionArguments functionArguments,
+ List<Type> inputDataTypes,
+ IoTDBLocal ioTDBLocal) {
+ checkArgument(ioTDBLocal != null, "IoTDBLocal must not be null for UDAF");
this.aggregateFunction = aggregateFunction;
+ this.functionArguments = functionArguments;
this.stateArray = new ObjectBigArray<>();
this.inputDataTypes = inputDataTypes;
+ this.ioTDBLocal = ioTDBLocal;
+ }
+
+ private void initIfNeeded() {
+ if (init) {
+ return;
+ }
+ try {
+ aggregateFunction.beforeStart(functionArguments, ioTDBLocal);
+ init = true;
+ } catch (UDFException e) {
+ throw new IoTDBRuntimeException(e, EXECUTE_UDF_ERROR.getStatusCode());
+ }
}
@Override
@@ -78,6 +104,7 @@ public class GroupedUserDefinedAggregateAccumulator implements GroupedAccumulato
@Override
public void addInput(int[] groupIds, Column[] arguments, AggregationMask mask) {
+ initIfNeeded();
RecordIterator iterator =
mask.isSelectAll()
? new RecordIterator(
@@ -90,7 +117,7 @@ public class GroupedUserDefinedAggregateAccumulator implements GroupedAccumulato
int groupId = groupIds[index];
index++;
State state = getOrCreateState(groupId);
- aggregateFunction.addInput(state, iterator.next());
+ aggregateFunction.addInput(state, iterator.next(), ioTDBLocal);
}
} else {
int[] selectedPositions = mask.getSelectedPositions();
@@ -98,13 +125,14 @@ public class GroupedUserDefinedAggregateAccumulator implements GroupedAccumulato
int groupId = groupIds[selectedPositions[index]];
index++;
State state = getOrCreateState(groupId);
- aggregateFunction.addInput(state, iterator.next());
+ aggregateFunction.addInput(state, iterator.next(), ioTDBLocal);
}
}
}
@Override
public void addIntermediate(int[] groupIds, Column argument) {
+ initIfNeeded();
checkArgument(
argument instanceof BinaryColumn
|| (argument instanceof RunLengthEncodedColumn
@@ -116,13 +144,14 @@ public class GroupedUserDefinedAggregateAccumulator implements GroupedAccumulato
State otherState = aggregateFunction.createState();
Binary otherStateBinary = argument.getBinary(i);
otherState.deserialize(otherStateBinary.getValues());
- aggregateFunction.combineState(getOrCreateState(groupIds[i]), otherState);
+ aggregateFunction.combineState(getOrCreateState(groupIds[i]), otherState, ioTDBLocal);
}
}
}
@Override
public void evaluateIntermediate(int groupId, ColumnBuilder columnBuilder) {
+ initIfNeeded();
checkArgument(
columnBuilder instanceof BinaryColumnBuilder,
"intermediate input and output of UDAF should be BinaryColumn");
@@ -136,8 +165,9 @@ public class GroupedUserDefinedAggregateAccumulator implements GroupedAccumulato
@Override
public void evaluateFinal(int groupId, ColumnBuilder columnBuilder) {
+ initIfNeeded();
ResultValue resultValue = new ResultValue(columnBuilder);
- aggregateFunction.outputFinal(getOrCreateState(groupId), resultValue);
+ aggregateFunction.outputFinal(getOrCreateState(groupId), resultValue, ioTDBLocal);
}
@Override
@@ -152,7 +182,9 @@ public class GroupedUserDefinedAggregateAccumulator implements GroupedAccumulato
@Override
public void close() {
- aggregateFunction.beforeDestroy();
+ initIfNeeded();
+ aggregateFunction.beforeDestroy(ioTDBLocal);
+ ioTDBLocal.close();
stateArray.forEach(
state -> {
if (state != null) {
diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/relational/ColumnTransformerBuilder.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/relational/ColumnTransformerBuilder.java
index a25188361af..7d3c1ab9bc5 100644
--- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/relational/ColumnTransformerBuilder.java
+++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/relational/ColumnTransformerBuilder.java
@@ -20,6 +20,7 @@
package org.apache.iotdb.calc.execution.relational;
import org.apache.iotdb.calc.i18n.CalcMessages;
+import org.apache.iotdb.calc.plan.planner.TableOperatorGenerator.IoTDBLocalFactory;
import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager;
import org.apache.iotdb.calc.plan.relational.metadata.ITypeMetadata;
import org.apache.iotdb.calc.transformation.dag.column.ColumnTransformer;
@@ -1489,11 +1490,10 @@ public class ColumnTransformerBuilder
.collect(Collectors.toList()),
Collections.emptyMap());
ScalarFunctionAnalysis analysis = scalarFunction.analyze(parameters);
- scalarFunction.beforeStart(parameters);
Type returnType =
UDFDataTypeTransformer.transformUDFDataTypeToReadType(analysis.getOutputDataType());
return new UserDefineScalarFunctionTransformer(
- returnType, scalarFunction, childrenColumnTransformer);
+ returnType, scalarFunction, childrenColumnTransformer, parameters, context);
}
}
throw new IllegalArgumentException(
@@ -1954,6 +1954,14 @@ public class ColumnTransformerBuilder
@SuppressWarnings("unused")
private final Optional<MemoryReservationManager> memoryReservationManager;
+ private final String fragmentInstanceId;
+
+ private final String outerGlobalQueryId;
+
+ private final long outerQueryDeadlineMs;
+
+ @Nullable private final IoTDBLocalFactory ioTDBLocalFactory;
+
public Context(
SessionInfo sessionInfo,
List<LeafColumnTransformer> leafList,
@@ -1966,6 +1974,40 @@ public class ColumnTransformerBuilder
ITableTypeProvider typeProvider,
ITypeMetadata metadata,
@Nullable MemoryReservationManager memoryReservationManager) {
+ this(
+ sessionInfo,
+ leafList,
+ inputLocations,
+ cache,
+ hasSeen,
+ commonTransformerList,
+ inputDataTypes,
+ originSize,
+ typeProvider,
+ metadata,
+ memoryReservationManager,
+ null,
+ null,
+ -1L,
+ null);
+ }
+
+ public Context(
+ SessionInfo sessionInfo,
+ List<LeafColumnTransformer> leafList,
+ Map<Symbol, List<InputLocation>> inputLocations,
+ Map<Expression, ColumnTransformer> cache,
+ Map<Expression, ColumnTransformer> hasSeen,
+ List<ColumnTransformer> commonTransformerList,
+ List<TSDataType> inputDataTypes,
+ int originSize,
+ ITableTypeProvider typeProvider,
+ ITypeMetadata metadata,
+ @Nullable MemoryReservationManager memoryReservationManager,
+ String fragmentInstanceId,
+ String outerGlobalQueryId,
+ long outerQueryDeadlineMs,
+ @Nullable IoTDBLocalFactory ioTDBLocalFactory) {
this.sessionInfo = sessionInfo;
this.leafList = leafList;
this.inputLocations = inputLocations;
@@ -1977,6 +2019,31 @@ public class ColumnTransformerBuilder
this.typeProvider = typeProvider;
this.metadata = metadata;
this.memoryReservationManager = Optional.ofNullable(memoryReservationManager);
+ this.fragmentInstanceId = fragmentInstanceId;
+ this.outerGlobalQueryId = outerGlobalQueryId;
+ this.outerQueryDeadlineMs = outerQueryDeadlineMs;
+ this.ioTDBLocalFactory = ioTDBLocalFactory;
+ }
+
+ public SessionInfo getSessionInfo() {
+ return sessionInfo;
+ }
+
+ public String getFragmentInstanceId() {
+ return fragmentInstanceId;
+ }
+
+ public String getOuterGlobalQueryId() {
+ return outerGlobalQueryId;
+ }
+
+ public long getOuterQueryDeadlineMs() {
+ return outerQueryDeadlineMs;
+ }
+
+ @Nullable
+ public IoTDBLocalFactory getIoTDBLocalFactory() {
+ return ioTDBLocalFactory;
}
public Type getType(SymbolReference symbolReference) {
diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/plan/planner/TableOperatorGenerator.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/plan/planner/TableOperatorGenerator.java
index bfac196461e..86786404bc4 100644
--- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/plan/planner/TableOperatorGenerator.java
+++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/plan/planner/TableOperatorGenerator.java
@@ -109,6 +109,7 @@ import org.apache.iotdb.calc.plan.relational.planner.CastToTimestampLiteralVisit
import org.apache.iotdb.calc.transformation.dag.column.ColumnTransformer;
import org.apache.iotdb.calc.transformation.dag.column.leaf.LeafColumnTransformer;
import org.apache.iotdb.calc.utils.datastructure.SortKey;
+import org.apache.iotdb.common.rpc.thrift.TAggregationType;
import org.apache.iotdb.commons.exception.SemanticException;
import org.apache.iotdb.commons.queryengine.common.SessionInfo;
import org.apache.iotdb.commons.queryengine.plan.analyze.ITableTypeProvider;
@@ -169,6 +170,7 @@ import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Expression;
import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Literal;
import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.SymbolReference;
import org.apache.iotdb.commons.queryengine.plan.relational.type.InternalTypeManager;
+import org.apache.iotdb.udf.api.IoTDBLocal;
import org.apache.iotdb.udf.api.relational.TableFunction;
import org.apache.iotdb.udf.api.relational.table.TableFunctionProcessorProvider;
@@ -314,6 +316,11 @@ public abstract class TableOperatorGenerator<
PlanNodeId planNodeId,
C context) {
+ String fragmentInstanceId = getFragmentInstanceId(context);
+ String outerGlobalQueryId = getQueryId(context);
+ long outerQueryDeadlineMs = getOuterQueryDeadlineMs(context);
+ IoTDBLocalFactory ioTDBLocalFactory = getIoTDBLocalFactory(context);
+
final List<TSDataType> filterOutputDataTypes = new ArrayList<>(inputDataTypes);
// records LeafColumnTransformer of filter
@@ -341,7 +348,11 @@ public abstract class TableOperatorGenerator<
0,
context.getTableTypeProvider(),
metadata,
- context.getMemoryReservationManager());
+ context.getMemoryReservationManager(),
+ fragmentInstanceId,
+ outerGlobalQueryId,
+ outerQueryDeadlineMs,
+ ioTDBLocalFactory);
return visitor.process(p, filterColumnTransformerContext);
})
@@ -369,7 +380,11 @@ public abstract class TableOperatorGenerator<
inputLocations.size(),
context.getTableTypeProvider(),
metadata,
- context.getMemoryReservationManager());
+ context.getMemoryReservationManager(),
+ fragmentInstanceId,
+ outerGlobalQueryId,
+ outerQueryDeadlineMs,
+ ioTDBLocalFactory);
for (Expression expression : projectExpressions) {
projectOutputTransformerList.add(
@@ -393,6 +408,22 @@ public abstract class TableOperatorGenerator<
predicate.isPresent());
}
+ protected String getFragmentInstanceId(C context) {
+ return null;
+ }
+
+ protected String getQueryId(C context) {
+ return null;
+ }
+
+ protected long getOuterQueryDeadlineMs(C context) {
+ return -1L;
+ }
+
+ protected IoTDBLocalFactory getIoTDBLocalFactory(C context) {
+ return null;
+ }
+
@Override
public Operator visitProject(ProjectNode node, C context) {
ITableTypeProvider typeProvider = context.getTableTypeProvider();
@@ -1329,7 +1360,8 @@ public abstract class TableOperatorGenerator<
false,
null,
Collections.emptySet(),
- operatorContext.getMemoryReservationContext())));
+ operatorContext.getMemoryReservationContext(),
+ context)));
return createAggregationOperator(operatorContext, child, aggregatorBuilder.build());
}
@@ -1344,7 +1376,8 @@ public abstract class TableOperatorGenerator<
boolean isAggTableScan,
String timeColumnName,
Set<String> measurementColumnNames,
- MemoryReservationManager memoryReservationManager) {
+ MemoryReservationManager memoryReservationManager,
+ C context) {
List<Integer> argumentChannels = new ArrayList<>();
for (Expression argument : aggregation.getArguments()) {
Symbol argumentSymbol = Symbol.from(argument);
@@ -1356,6 +1389,10 @@ public abstract class TableOperatorGenerator<
aggregation.getResolvedFunction().getSignature().getArgumentTypes().stream()
.map(InternalTypeManager::getTSDataType)
.collect(Collectors.toList());
+ IoTDBLocal ioTDBLocal =
+ getAggregationTypeByFuncName(functionName) == TAggregationType.UDAF
+ ? createIoTDBLocal(context)
+ : null;
TableAccumulator accumulator =
createAccumulator(
functionName,
@@ -1368,7 +1405,8 @@ public abstract class TableOperatorGenerator<
timeColumnName,
measurementColumnNames,
aggregation.isDistinct(),
- memoryReservationManager);
+ memoryReservationManager,
+ ioTDBLocal);
OptionalInt maskChannel = OptionalInt.empty();
if (aggregation.hasMask()) {
@@ -1411,7 +1449,8 @@ public abstract class TableOperatorGenerator<
false,
null,
Collections.emptySet(),
- context.getMemoryReservationManager())));
+ context.getMemoryReservationManager(),
+ context)));
CommonOperatorContext operatorContext =
addOperatorContext(
@@ -1439,7 +1478,8 @@ public abstract class TableOperatorGenerator<
v,
node.getStep(),
typeProvider,
- context.getMemoryReservationManager())));
+ context.getMemoryReservationManager(),
+ context)));
Set<Symbol> preGroupedKeys = ImmutableSet.copyOf(node.getPreGroupedSymbols());
List<Symbol> groupingKeys = node.getGroupingKeys();
@@ -1496,7 +1536,8 @@ public abstract class TableOperatorGenerator<
v,
node.getStep(),
typeProvider,
- context.getMemoryReservationManager())));
+ context.getMemoryReservationManager(),
+ context)));
CommonOperatorContext operatorContext =
addOperatorContext(
context, node.getPlanNodeId(), HashAggregationOperator.class.getSimpleName());
@@ -1621,7 +1662,8 @@ public abstract class TableOperatorGenerator<
AggregationNode.Aggregation aggregation,
AggregationNode.Step step,
ITableTypeProvider typeProvider,
- MemoryReservationManager memoryReservationManager) {
+ MemoryReservationManager memoryReservationManager,
+ C context) {
List<Integer> argumentChannels = new ArrayList<>();
for (Expression argument : aggregation.getArguments()) {
Symbol argumentSymbol = Symbol.from(argument);
@@ -1633,6 +1675,10 @@ public abstract class TableOperatorGenerator<
aggregation.getResolvedFunction().getSignature().getArgumentTypes().stream()
.map(InternalTypeManager::getTSDataType)
.collect(Collectors.toList());
+ IoTDBLocal ioTDBLocal =
+ getAggregationTypeByFuncName(functionName) == TAggregationType.UDAF
+ ? createIoTDBLocal(context)
+ : null;
GroupedAccumulator accumulator =
createGroupedAccumulator(
functionName,
@@ -1642,7 +1688,8 @@ public abstract class TableOperatorGenerator<
Collections.emptyMap(),
true,
aggregation.isDistinct(),
- memoryReservationManager);
+ memoryReservationManager,
+ ioTDBLocal);
OptionalInt maskChannel = OptionalInt.empty();
if (aggregation.hasMask()) {
@@ -1671,7 +1718,9 @@ public abstract class TableOperatorGenerator<
CommonOperatorContext operatorContext =
addOperatorContext(
context, node.getPlanNodeId(), TableFunctionLeafOperator.class.getSimpleName());
- return new TableFunctionLeafOperator(operatorContext, processorProvider, outputDataTypes);
+ IoTDBLocal ioTDBLocal = createIoTDBLocal(context);
+ return new TableFunctionLeafOperator(
+ operatorContext, processorProvider, outputDataTypes, ioTDBLocal);
} else {
Operator operator = node.getChild().accept(this, context);
CommonOperatorContext operatorContext =
@@ -1716,6 +1765,7 @@ public abstract class TableOperatorGenerator<
} else {
partitionChannels = Collections.emptyList();
}
+ IoTDBLocal ioTDBLocal = createIoTDBLocal(context);
return new TableFunctionOperator(
operatorContext,
processorProvider,
@@ -1729,7 +1779,8 @@ public abstract class TableOperatorGenerator<
.map(TableFunctionNode.PassThroughSpecification::isDeclaredAsPassThrough)
.orElse(false),
partitionChannels,
- node.isRequireRecordSnapshot());
+ node.isRequireRecordSnapshot(),
+ ioTDBLocal);
}
}
@@ -2467,4 +2518,40 @@ public abstract class TableOperatorGenerator<
}
protected abstract SessionInfo getSessionInfo(C context);
+
+ protected IoTDBLocal createIoTDBLocal(C context) {
+ return IoTDBLocalFactory.createIoTDBLocal(
+ getIoTDBLocalFactory(context),
+ getSessionInfo(context),
+ getFragmentInstanceId(context),
+ getQueryId(context),
+ getOuterQueryDeadlineMs(context));
+ }
+
+ /** Factory for creating {@link IoTDBLocal} inside UDF column transformers. */
+ @FunctionalInterface
+ public interface IoTDBLocalFactory {
+
+ IoTDBLocal create(
+ SessionInfo sessionInfo,
+ String fragmentInstanceId,
+ String outerGlobalQueryId,
+ long outerQueryDeadlineMs);
+
+ static IoTDBLocal createIoTDBLocal(
+ IoTDBLocalFactory factory,
+ SessionInfo sessionInfo,
+ String fragmentInstanceId,
+ String outerGlobalQueryId,
+ long outerQueryDeadlineMs) {
+ checkArgument(factory != null, "IoTDBLocalFactory must not be null for UDF execution");
+ checkArgument(
+ fragmentInstanceId != null, "fragmentInstanceId must not be null for UDF execution");
+ checkArgument(outerGlobalQueryId != null, "queryId must not be null for UDF execution");
+ checkArgument(
+ outerQueryDeadlineMs > 0, "outerQueryDeadlineMs must be positive for UDF execution");
+ return factory.create(
+ sessionInfo, fragmentInstanceId, outerGlobalQueryId, outerQueryDeadlineMs);
+ }
+ }
}
diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/transformation/dag/column/udf/UserDefineScalarFunctionTransformer.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/transformation/dag/column/udf/UserDefineScalarFunctionTransformer.java
index 4e22315b628..3288dcbad07 100644
--- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/transformation/dag/column/udf/UserDefineScalarFunctionTransformer.java
+++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/transformation/dag/column/udf/UserDefineScalarFunctionTransformer.java
@@ -20,8 +20,14 @@
package org.apache.iotdb.calc.transformation.dag.column.udf;
import org.apache.iotdb.calc.execution.operator.source.relational.aggregation.RecordIterator;
+import org.apache.iotdb.calc.execution.relational.ColumnTransformerBuilder;
+import org.apache.iotdb.calc.plan.planner.TableOperatorGenerator.IoTDBLocalFactory;
import org.apache.iotdb.calc.transformation.dag.column.ColumnTransformer;
import org.apache.iotdb.calc.transformation.dag.column.multi.MultiColumnTransformer;
+import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
+import org.apache.iotdb.udf.api.IoTDBLocal;
+import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;
+import org.apache.iotdb.udf.api.exception.UDFException;
import org.apache.iotdb.udf.api.relational.ScalarFunction;
import org.apache.iotdb.udf.api.relational.access.Record;
@@ -32,28 +38,56 @@ import org.apache.tsfile.read.common.type.Type;
import java.util.List;
import java.util.stream.Collectors;
+import static org.apache.iotdb.rpc.TSStatusCode.EXECUTE_UDF_ERROR;
+
public class UserDefineScalarFunctionTransformer extends MultiColumnTransformer {
private final ScalarFunction scalarFunction;
+ private final FunctionArguments parameters;
private final List<Type> inputTypes;
+ private final IoTDBLocal ioTDBLocal;
+ private boolean init = false;
public UserDefineScalarFunctionTransformer(
Type returnType,
ScalarFunction scalarFunction,
- List<ColumnTransformer> childrenTransformers) {
+ List<ColumnTransformer> childrenTransformers,
+ FunctionArguments parameters,
+ ColumnTransformerBuilder.Context context) {
super(returnType, childrenTransformers);
this.scalarFunction = scalarFunction;
+ this.parameters = parameters;
+ this.ioTDBLocal =
+ IoTDBLocalFactory.createIoTDBLocal(
+ context.getIoTDBLocalFactory(),
+ context.getSessionInfo(),
+ context.getFragmentInstanceId(),
+ context.getOuterGlobalQueryId(),
+ context.getOuterQueryDeadlineMs());
this.inputTypes =
childrenTransformers.stream().map(ColumnTransformer::getType).collect(Collectors.toList());
}
+ private void initIfNeeded() {
+ if (init) {
+ return;
+ }
+ init = true;
+ try {
+ scalarFunction.beforeStart(parameters, ioTDBLocal);
+ } catch (UDFException e) {
+ throw new IoTDBRuntimeException(e, EXECUTE_UDF_ERROR.getStatusCode());
+ }
+ }
+
@Override
protected void doTransform(
List<Column> childrenColumns, ColumnBuilder builder, int positionCount) {
+ initIfNeeded();
RecordIterator iterator = new RecordIterator(childrenColumns, inputTypes, positionCount);
while (iterator.hasNext()) {
try {
- Object result = scalarFunction.evaluate(iterator.next());
+ Object result = scalarFunction.evaluate(iterator.next(), ioTDBLocal);
if (result == null) {
builder.appendNull();
} else {
@@ -71,6 +105,7 @@ public class UserDefineScalarFunctionTransformer extends MultiColumnTransformer
@Override
protected void doTransform(
List<Column> childrenColumns, ColumnBuilder builder, int positionCount, boolean[] selection) {
+ initIfNeeded();
RecordIterator iterator = new RecordIterator(childrenColumns, inputTypes, positionCount);
int i = 0;
while (iterator.hasNext()) {
@@ -80,7 +115,7 @@ public class UserDefineScalarFunctionTransformer extends MultiColumnTransformer
builder.appendNull();
continue;
}
- Object result = scalarFunction.evaluate(input);
+ Object result = scalarFunction.evaluate(input, ioTDBLocal);
if (result == null) {
builder.appendNull();
} else {
@@ -97,8 +132,11 @@ public class UserDefineScalarFunctionTransformer extends MultiColumnTransformer
@Override
public void close() {
+ // ensure beforeStart was called
+ initIfNeeded();
super.close();
- scalarFunction.beforeDestroy();
+ scalarFunction.beforeDestroy(ioTDBLocal);
+ ioTDBLocal.close();
}
@Override
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java
index b3fa3c049f3..21dbffb606e 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java
@@ -369,6 +369,31 @@ public class SessionManager implements SessionManagerMBean {
return currSession.get();
}
+ /**
+ * Temporarily install a session into the current thread. Used by UDF internal queries that need
+ * ThreadLocal session visibility without {@link #registerSession(IClientSession)}.
+ */
+ public void setCurrSession(IClientSession session) {
+ currSession.set(session);
+ sessions.putIfAbsent(session, placeHolder);
+ }
+
+ /**
+ * Restore the previous ThreadLocal session and remove the temporarily installed session from
+ * {@code sessions}.
+ */
+ public void restoreSession(IClientSession previous, IClientSession installedSession) {
+ if (installedSession != null) {
+ sessions.remove(installedSession);
+ }
+ if (previous != null) {
+ currSession.set(previous);
+ } else {
+ currSession.remove();
+ currSessionIdleTime.remove();
+ }
+ }
+
/** get current session and update session idle time. */
public IClientSession getCurrSessionAndUpdateIdleTime() {
IClientSession clientSession = getCurrSession();
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java
index 6b0350287fb..50312a4cea6 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java
@@ -562,7 +562,7 @@ public class ClientRPCServiceImpl implements IClientRPCServiceWithHandler {
}
}
- private void clearUp(
+ public static void clearUp(
IClientSession clientSession,
Long statementId,
Long queryId,
@@ -572,7 +572,7 @@ public class ClientRPCServiceImpl implements IClientRPCServiceWithHandler {
clientSession.removeQueryId(statementId, queryId);
}
- private void clearUp(
+ private static void clearUp(
IClientSession clientSession,
Long statementId,
Long queryId,
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java
index 3ffa7e63e15..81dd89387bf 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java
@@ -161,6 +161,9 @@ public class FragmentInstanceContext extends QueryContext {
// session info
private SessionInfo sessionInfo;
+ // Outer query deadline (startTime + timeout) for IoTDBLocal UDF
+ private long outerQueryDeadlineMs = -1L;
+
private final Map<QueryId, DataNodeQueryContext> dataNodeQueryContextMap;
private DataNodeQueryContext dataNodeQueryContext;
@@ -218,6 +221,7 @@ public class FragmentInstanceContext extends QueryContext {
IDataRegionForQuery dataRegion,
TimePredicate globalTimePredicate,
Map<QueryId, DataNodeQueryContext> dataNodeQueryContextMap,
+ long outerQueryDeadlineMs,
boolean debug,
boolean isVerbose) {
FragmentInstanceContext instanceContext =
@@ -228,6 +232,7 @@ public class FragmentInstanceContext extends QueryContext {
dataRegion,
globalTimePredicate,
dataNodeQueryContextMap,
+ outerQueryDeadlineMs,
debug,
isVerbose);
instanceContext.initialize();
@@ -288,6 +293,7 @@ public class FragmentInstanceContext extends QueryContext {
IDataRegionForQuery dataRegion,
TimePredicate globalTimePredicate,
Map<QueryId, DataNodeQueryContext> dataNodeQueryContextMap,
+ long outerQueryDeadlineMs,
boolean debug,
boolean verbose) {
super(debug, verbose);
@@ -295,6 +301,7 @@ public class FragmentInstanceContext extends QueryContext {
this.stateMachine = stateMachine;
this.executionEndTime.set(END_TIME_INITIAL_VALUE);
this.sessionInfo = sessionInfo;
+ this.outerQueryDeadlineMs = outerQueryDeadlineMs;
this.dataRegion = dataRegion;
this.globalTimeFilter =
globalTimePredicate == null
@@ -584,6 +591,10 @@ public class FragmentInstanceContext extends QueryContext {
return sessionInfo;
}
+ public long getOuterQueryDeadlineMs() {
+ return outerQueryDeadlineMs;
+ }
+
public Optional<Throwable> getFailureCause() {
return Optional.ofNullable(
stateMachine.getFailureCauses().stream()
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java
index e5f3cbc7e0d..a9473982237 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java
@@ -38,6 +38,7 @@ import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeService
import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISink;
import org.apache.iotdb.db.queryengine.execution.schedule.DriverScheduler;
import org.apache.iotdb.db.queryengine.execution.schedule.IDriverScheduler;
+import org.apache.iotdb.db.queryengine.execution.schedule.task.DriverTask;
import org.apache.iotdb.db.queryengine.metric.QueryRelatedResourceMetricSet;
import org.apache.iotdb.db.queryengine.plan.Coordinator;
import org.apache.iotdb.db.queryengine.plan.planner.LocalExecutionPlanner;
@@ -167,6 +168,7 @@ public class FragmentInstanceManager {
dataRegion,
instance.getGlobalTimePredicate(),
dataNodeQueryContextMap,
+ DriverTask.computeDeadlineTimeInMs(instance.getTimeOut()),
instance.isDebug(),
instance.isVerbose());
});
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/task/DriverTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/task/DriverTask.java
index c85d00e17e1..285c67dd7ba 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/task/DriverTask.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/task/DriverTask.java
@@ -79,15 +79,7 @@ public class DriverTask implements IDIndexedAccessible {
this.driver = driver;
this.setStatus(status);
- long currentTime = System.currentTimeMillis();
- long ddlTmp = currentTime + timeoutMs;
- // avoid infinite timeout check loop, schema fetch query for write operation may pass a very
- // large timeout here which may causing currentTime + timeoutMs be negative
- if (ddlTmp < currentTime) {
- this.ddl = Long.MAX_VALUE;
- } else {
- this.ddl = ddlTmp;
- }
+ this.ddl = computeDeadlineTimeInMs(timeoutMs);
this.lock = new ReentrantLock();
this.driverTaskHandle = driverTaskHandle;
@@ -96,6 +88,17 @@ public class DriverTask implements IDIndexedAccessible {
this.isHighestPriority = isHighestPriority;
}
+ public static long computeDeadlineTimeInMs(long timeoutMs) {
+ long currentTime = System.currentTimeMillis();
+ long deadlineMs = currentTime + timeoutMs;
+ // avoid infinite timeout check loop, schema fetch query for write operation may pass a very
+ // large timeout here which may causing currentTime + timeoutMs be negative
+ if (deadlineMs < currentTime) {
+ return Long.MAX_VALUE;
+ }
+ return deadlineMs;
+ }
+
@Override
public DriverTaskId getDriverTaskId() {
return driver.getDriverTaskId();
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java
index 4a552ddbeb0..c0daee843ce 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java
@@ -156,6 +156,7 @@ import org.apache.iotdb.db.queryengine.plan.relational.sql.rewrite.StatementRewr
import org.apache.iotdb.db.queryengine.plan.statement.IConfigStatement;
import org.apache.iotdb.db.queryengine.plan.statement.Statement;
import org.apache.iotdb.db.queryengine.plan.statement.pipe.PipeEnrichedStatement;
+import org.apache.iotdb.db.queryengine.udf.InternalQueryExecutor;
import org.apache.iotdb.db.utils.SetThreadName;
import org.apache.thrift.TBase;
@@ -308,6 +309,17 @@ public class Coordinator {
boolean userQuery,
boolean debug,
BiFunction<MPPQueryContext, Long, IQueryExecution> iQueryExecutionFactory) {
+ return execution(queryId, session, sql, userQuery, debug, false, iQueryExecutionFactory);
+ }
+
+ private ExecutionResult execution(
+ long queryId,
+ SessionInfo session,
+ String sql,
+ boolean userQuery,
+ boolean debug,
+ boolean readOnlyInternalQuery,
+ BiFunction<MPPQueryContext, Long, IQueryExecution> iQueryExecutionFactory) {
long startTime = System.currentTimeMillis();
QueryId globalQueryId = queryIdGenerator.createNextQueryId();
MPPQueryContext queryContext = null;
@@ -326,6 +338,9 @@ public class Coordinator {
queryContext.setUserQuery(userQuery);
queryContext.setDebug(debug);
IQueryExecution execution = iQueryExecutionFactory.apply(queryContext, startTime);
+ if (readOnlyInternalQuery) {
+ InternalQueryExecutor.validateReadOnlyQuery(execution);
+ }
if (execution.isQuery()) {
queryExecutionMap.put(queryId, execution);
} else {
@@ -489,12 +504,39 @@ public class Coordinator {
long timeOut,
boolean userQuery,
boolean debug) {
+ return executeForTableModel(
+ statement,
+ sqlParser,
+ clientSession,
+ queryId,
+ session,
+ sql,
+ metadata,
+ timeOut,
+ userQuery,
+ debug,
+ false);
+ }
+
+ public ExecutionResult executeForTableModel(
+ org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Statement statement,
+ SqlParser sqlParser,
+ IClientSession clientSession,
+ long queryId,
+ SessionInfo session,
+ String sql,
+ Metadata metadata,
+ long timeOut,
+ boolean userQuery,
+ boolean debug,
+ boolean readOnlyInternalQuery) {
return execution(
queryId,
session,
sql,
userQuery,
debug,
+ readOnlyInternalQuery,
((queryContext, startTime) ->
createQueryExecutionForTableModel(
statement,
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java
index aa5b7992303..37dbd3ba741 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java
@@ -28,6 +28,7 @@ import org.apache.iotdb.calc.execution.operator.source.relational.aggregation.La
import org.apache.iotdb.calc.execution.operator.source.relational.aggregation.TableAggregator;
import org.apache.iotdb.calc.execution.relational.ColumnTransformerBuilder;
import org.apache.iotdb.calc.plan.planner.TableOperatorGenerator;
+import org.apache.iotdb.calc.plan.planner.TableOperatorGenerator.IoTDBLocalFactory;
import org.apache.iotdb.calc.transformation.dag.column.leaf.LeafColumnTransformer;
import org.apache.iotdb.calc.transformation.dag.column.unary.scalar.DateBinFunctionColumnTransformer;
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
@@ -136,6 +137,7 @@ import org.apache.iotdb.db.queryengine.plan.relational.planner.node.schema.Table
import org.apache.iotdb.db.queryengine.plan.relational.planner.node.schema.TableDeviceQueryCountNode;
import org.apache.iotdb.db.queryengine.plan.relational.planner.node.schema.TableDeviceQueryScanNode;
import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering;
+import org.apache.iotdb.db.queryengine.udf.IoTDBLocalImpl;
import org.apache.iotdb.db.schemaengine.schemaregion.read.resp.info.IDeviceSchemaInfo;
import org.apache.iotdb.db.schemaengine.table.DataNodeTableCache;
import org.apache.iotdb.db.schemaengine.table.DataNodeTreeViewSchemaUtils;
@@ -1503,7 +1505,8 @@ public class DataNodeTableOperatorGenerator
true,
timeColumnName,
measurementColumnsIndexMap.keySet(),
- context.getMemoryReservationManager()));
+ context.getMemoryReservationManager(),
+ context));
}
ITableTimeRangeIterator timeRangeIterator = null;
@@ -2168,4 +2171,24 @@ public class DataNodeTableOperatorGenerator
protected SessionInfo getSessionInfo(LocalExecutionPlanContext context) {
return context.getDriverContext().getFragmentInstanceContext().getSessionInfo();
}
+
+ @Override
+ protected String getFragmentInstanceId(LocalExecutionPlanContext context) {
+ return context.getFragmentInstanceId().getFullId();
+ }
+
+ @Override
+ protected String getQueryId(LocalExecutionPlanContext context) {
+ return context.getFragmentInstanceId().getQueryId().getId();
+ }
+
+ @Override
+ protected long getOuterQueryDeadlineMs(LocalExecutionPlanContext context) {
+ return context.getOuterQueryDeadlineMs();
+ }
+
+ @Override
+ protected IoTDBLocalFactory getIoTDBLocalFactory(LocalExecutionPlanContext context) {
+ return IoTDBLocalImpl.FACTORY;
+ }
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java
index 65908708394..4d78236f38d 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java
@@ -198,6 +198,10 @@ public class LocalExecutionPlanContext implements ITableOperatorGeneratorContext
return driverContext.getFragmentInstanceContext().getId();
}
+ public long getOuterQueryDeadlineMs() {
+ return driverContext.getFragmentInstanceContext().getOuterQueryDeadlineMs();
+ }
+
public List<PipelineDriverFactory> getPipelineDriverFactories() {
return pipelineDriverFactories;
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/udf/InternalQueryExecutor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/udf/InternalQueryExecutor.java
new file mode 100644
index 00000000000..5abcf05d22b
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/udf/InternalQueryExecutor.java
@@ -0,0 +1,133 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.queryengine.udf;
+
+import org.apache.iotdb.commons.exception.IoTDBException;
+import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
+import org.apache.iotdb.commons.exception.SemanticException;
+import org.apache.iotdb.commons.queryengine.common.SessionInfo;
+import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Statement;
+import org.apache.iotdb.db.protocol.session.IClientSession;
+import org.apache.iotdb.db.protocol.session.InternalClientSession;
+import org.apache.iotdb.db.protocol.session.SessionManager;
+import org.apache.iotdb.db.protocol.thrift.impl.ClientRPCServiceImpl;
+import org.apache.iotdb.db.queryengine.common.QueryId;
+import org.apache.iotdb.db.queryengine.plan.Coordinator;
+import org.apache.iotdb.db.queryengine.plan.analyze.QueryType;
+import org.apache.iotdb.db.queryengine.plan.execution.ExecutionResult;
+import org.apache.iotdb.db.queryengine.plan.execution.IQueryExecution;
+import org.apache.iotdb.db.queryengine.plan.planner.LocalExecutionPlanner;
+import org.apache.iotdb.db.queryengine.plan.relational.metadata.Metadata;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.parser.SqlParser;
+import org.apache.iotdb.rpc.TSStatusCode;
+
+/** Stateless utility for UDF embedded table-model queries. */
+public final class InternalQueryExecutor {
+
+ private static final Coordinator COORDINATOR = Coordinator.getInstance();
+ private static final SessionManager SESSION_MANAGER = SessionManager.getInstance();
+ private static final Metadata METADATA = LocalExecutionPlanner.getInstance().metadata;
+ private static final SqlParser RELATION_SQL_PARSER = new SqlParser();
+
+ private InternalQueryExecutor() {}
+
+ public static InternalQueryResult executeInternalQuery(
+ SessionInfo sessionInfo,
+ String fragmentInstanceId,
+ QueryId outerQueryId,
+ String sql,
+ long timeoutMs)
+ throws IoTDBException {
+
+ IClientSession previousSession = SESSION_MANAGER.getCurrSession();
+
+ InternalClientSession internalSession =
+ new InternalClientSession(formatInternalClientId(fragmentInstanceId, outerQueryId));
+ internalSession.setSqlDialect(sessionInfo.getSqlDialect());
+ sessionInfo.getDatabaseName().ifPresent(internalSession::setDatabaseName);
+
+ SESSION_MANAGER.supplySession(
+ internalSession,
+ sessionInfo.getUserId(),
+ sessionInfo.getUserName(),
+ sessionInfo.getZoneId(),
+ sessionInfo.getVersion());
+
+ long statementId = -1;
+ long queryId = -1;
+ try {
+ SESSION_MANAGER.setCurrSession(internalSession);
+
+ statementId = SESSION_MANAGER.requestStatementId(internalSession);
+ queryId = SESSION_MANAGER.requestQueryId(internalSession, statementId);
+
+ Statement parsedStatement =
+ RELATION_SQL_PARSER.createStatement(sql, sessionInfo.getZoneId(), internalSession);
+
+ ExecutionResult result =
+ COORDINATOR.executeForTableModel(
+ parsedStatement,
+ RELATION_SQL_PARSER,
+ internalSession,
+ queryId,
+ sessionInfo,
+ sql,
+ METADATA,
+ timeoutMs,
+ false,
+ false,
+ true);
+
+ if (result.status.code != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+ throw new IoTDBException(result.status.message, result.status.code);
+ }
+
+ IQueryExecution queryExecution = COORDINATOR.getQueryExecution(queryId);
+ if (queryExecution == null) {
+ throw new IoTDBException(
+ "Internal query execution not found",
+ TSStatusCode.INTERNAL_SERVER_ERROR.getStatusCode());
+ }
+
+ return new InternalQueryResult(queryExecution, internalSession, statementId, queryId, sql);
+ } catch (IoTDBException e) {
+ ClientRPCServiceImpl.clearUp(internalSession, statementId, queryId, () -> sql, e);
+ throw e;
+ } catch (IoTDBRuntimeException e) {
+ ClientRPCServiceImpl.clearUp(internalSession, statementId, queryId, () -> sql, e);
+ throw e;
+ } catch (Exception e) {
+ ClientRPCServiceImpl.clearUp(internalSession, statementId, queryId, () -> sql, e);
+ throw new IoTDBException(e.getMessage(), TSStatusCode.INTERNAL_SERVER_ERROR.getStatusCode());
+ } finally {
+ SESSION_MANAGER.restoreSession(previousSession, internalSession);
+ }
+ }
+
+ static String formatInternalClientId(String fragmentInstanceId, QueryId outerQueryId) {
+ return String.format("udf-local-%s-%s", fragmentInstanceId, outerQueryId);
+ }
+
+ public static void validateReadOnlyQuery(IQueryExecution execution) {
+ if (execution.getQueryType() != QueryType.READ) {
+ throw new SemanticException("Only query is supported for IoTDBLocal query interface");
+ }
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/udf/InternalQueryResult.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/udf/InternalQueryResult.java
new file mode 100644
index 00000000000..419ea00eb00
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/udf/InternalQueryResult.java
@@ -0,0 +1,65 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.queryengine.udf;
+
+import org.apache.iotdb.db.protocol.session.IClientSession;
+import org.apache.iotdb.db.protocol.thrift.impl.ClientRPCServiceImpl;
+import org.apache.iotdb.db.queryengine.common.header.DatasetHeader;
+import org.apache.iotdb.db.queryengine.plan.execution.IQueryExecution;
+
+/** Internal query result holding {@link IQueryExecution} and cleanup metadata. */
+public final class InternalQueryResult implements AutoCloseable {
+
+ private final IQueryExecution queryExecution;
+ private final IClientSession internalSession;
+ private final long statementId;
+ private final long queryId;
+ private final String sql;
+
+ public InternalQueryResult(
+ IQueryExecution queryExecution,
+ IClientSession internalSession,
+ long statementId,
+ long queryId,
+ String sql) {
+ this.queryExecution = queryExecution;
+ this.internalSession = internalSession;
+ this.statementId = statementId;
+ this.queryId = queryId;
+ this.sql = sql;
+ }
+
+ public IQueryExecution getQueryExecution() {
+ return queryExecution;
+ }
+
+ public long getQueryId() {
+ return queryId;
+ }
+
+ public DatasetHeader getDatasetHeader() {
+ return queryExecution.getDatasetHeader();
+ }
+
+ @Override
+ public void close() {
+ ClientRPCServiceImpl.clearUp(internalSession, statementId, queryId, () -> sql, null);
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/udf/IoTDBLocalImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/udf/IoTDBLocalImpl.java
new file mode 100644
index 00000000000..630360f4858
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/udf/IoTDBLocalImpl.java
@@ -0,0 +1,147 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.queryengine.udf;
+
+import org.apache.iotdb.calc.plan.planner.TableOperatorGenerator.IoTDBLocalFactory;
+import org.apache.iotdb.commons.exception.IoTDBException;
+import org.apache.iotdb.commons.exception.QueryTimeoutException;
+import org.apache.iotdb.commons.queryengine.common.SessionInfo;
+import org.apache.iotdb.db.queryengine.common.QueryId;
+import org.apache.iotdb.udf.api.IoTDBLocal;
+import org.apache.iotdb.udf.api.UDFResultSet;
+import org.apache.iotdb.udf.api.exception.UDFException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/** Server-side implementation of {@link IoTDBLocal}. */
+public class IoTDBLocalImpl implements IoTDBLocal {
+
+ public static final IoTDBLocalFactory FACTORY = IoTDBLocalImpl::new;
+
+ private static final Logger LOGGER = LoggerFactory.getLogger(IoTDBLocalImpl.class);
+
+ private final SessionInfo sessionInfo;
+ private final String fragmentInstanceId;
+ private final QueryId outerGlobalQueryId;
+ private final long outerQueryDeadlineMs;
+ private final List<UDFResultSetImpl> openResultSets = new ArrayList<>();
+
+ public IoTDBLocalImpl(
+ SessionInfo sessionInfo,
+ String fragmentInstanceId,
+ String outerGlobalQueryId,
+ long outerQueryDeadlineMs) {
+ this.sessionInfo = sessionInfo;
+ this.fragmentInstanceId = fragmentInstanceId;
+ this.outerGlobalQueryId = QueryId.valueOf(outerGlobalQueryId);
+ this.outerQueryDeadlineMs = outerQueryDeadlineMs;
+ }
+
+ @Override
+ public UDFResultSet query(String sql) throws UDFException {
+ try {
+ long timeoutMs = computeRemainingTimeoutMs();
+ if (timeoutMs <= 0) {
+ throw new QueryTimeoutException(
+ "Outer query timeout exceeded before IoTDBLocal query starts");
+ }
+ InternalQueryResult result =
+ InternalQueryExecutor.executeInternalQuery(
+ sessionInfo, fragmentInstanceId, outerGlobalQueryId, sql, timeoutMs);
+ int index = openResultSets.size();
+ UDFResultSetImpl rs = new UDFResultSetImpl(openResultSets, index, result);
+ openResultSets.add(rs);
+ return rs;
+ } catch (IoTDBException e) {
+ throw new UDFException(e.getMessage(), e);
+ }
+ }
+
+ @Override
+ public void close() {
+ for (int i = 0; i < openResultSets.size(); i++) {
+ UDFResultSetImpl rs = openResultSets.get(i);
+ if (rs != null) {
+ try {
+ rs.close();
+ } catch (UDFException e) {
+ LOGGER.warn("Failed to close UDF result set at index {}", i, e);
+ }
+ }
+ }
+ openResultSets.clear();
+ }
+
+ private long computeRemainingTimeoutMs() {
+ if (outerQueryDeadlineMs <= 0) {
+ return 0;
+ }
+ return outerQueryDeadlineMs - System.currentTimeMillis();
+ }
+
+ @Override
+ public void info(String msg) {
+ LOGGER.info(msg);
+ }
+
+ @Override
+ public void info(String format, Object... args) {
+ LOGGER.info(format, args);
+ }
+
+ @Override
+ public void info(String msg, Throwable t) {
+ LOGGER.info(msg, t);
+ }
+
+ @Override
+ public void warn(String msg) {
+ LOGGER.warn(msg);
+ }
+
+ @Override
+ public void warn(String format, Object... args) {
+ LOGGER.warn(format, args);
+ }
+
+ @Override
+ public void warn(String msg, Throwable t) {
+ LOGGER.warn(msg, t);
+ }
+
+ @Override
+ public void error(String msg) {
+ LOGGER.error(msg);
+ }
+
+ @Override
+ public void error(String format, Object... args) {
+ LOGGER.error(format, args);
+ }
+
+ @Override
+ public void error(String msg, Throwable t) {
+ LOGGER.error(msg, t);
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/udf/UDFResultSetImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/udf/UDFResultSetImpl.java
new file mode 100644
index 00000000000..74cc1a3859f
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/udf/UDFResultSetImpl.java
@@ -0,0 +1,154 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.queryengine.udf;
+
+import org.apache.iotdb.calc.execution.operator.source.relational.aggregation.RecordIterator;
+import org.apache.iotdb.commons.exception.IoTDBException;
+import org.apache.iotdb.commons.schema.column.ColumnHeader;
+import org.apache.iotdb.commons.udf.utils.UDFDataTypeTransformer;
+import org.apache.iotdb.db.queryengine.common.header.DatasetHeader;
+import org.apache.iotdb.udf.api.UDFResultSet;
+import org.apache.iotdb.udf.api.exception.UDFException;
+import org.apache.iotdb.udf.api.relational.access.Record;
+
+import org.apache.tsfile.block.column.Column;
+import org.apache.tsfile.read.common.block.TsBlock;
+import org.apache.tsfile.read.common.type.Type;
+
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.NoSuchElementException;
+import java.util.Optional;
+import java.util.stream.Collectors;
+
+/** Server-side implementation of {@link UDFResultSet}. */
+public class UDFResultSetImpl implements UDFResultSet {
+
+ private final List<UDFResultSetImpl> openResultSets;
+ private final int index;
+ private final InternalQueryResult queryResult;
+ private final List<Type> columnTypes;
+
+ /** null when output columns align with TsBlock value columns without index mapping */
+ private final int[] columnIndexes;
+
+ private Iterator<Record> rowIterator;
+ private boolean closed;
+
+ public UDFResultSetImpl(
+ List<UDFResultSetImpl> openResultSets, int index, InternalQueryResult queryResult) {
+ this.openResultSets = openResultSets;
+ this.index = index;
+ this.queryResult = queryResult;
+ DatasetHeader datasetHeader = queryResult.getDatasetHeader();
+ this.columnTypes = buildColumnTypes(datasetHeader.getColumnHeaders());
+ this.columnIndexes = buildColumnIndexes(datasetHeader.getColumnIndex2TsBlockColumnIndexList());
+ }
+
+ @Override
+ public boolean hasNext() throws UDFException {
+ ensureOpen();
+
+ while (rowIterator == null || !rowIterator.hasNext()) {
+ if (!queryResult.getQueryExecution().hasNextResult()) {
+ return false;
+ }
+ Optional<TsBlock> batch;
+ try {
+ batch = queryResult.getQueryExecution().getBatchResult();
+ } catch (IoTDBException e) {
+ throw new UDFException(e.getMessage(), e);
+ }
+ if (!batch.isPresent()) {
+ return false;
+ }
+ TsBlock currentBlock = batch.get();
+ if (currentBlock.getPositionCount() == 0) {
+ continue;
+ }
+ rowIterator =
+ new RecordIterator(
+ extractColumns(currentBlock), columnTypes, currentBlock.getPositionCount());
+ }
+ return true;
+ }
+
+ @Override
+ public Record next() throws UDFException {
+ if (!hasNext()) {
+ throw new NoSuchElementException();
+ }
+ return rowIterator.next();
+ }
+
+ @Override
+ public void close() throws UDFException {
+ if (closed) {
+ return;
+ }
+ closed = true;
+ openResultSets.set(index, null);
+ try {
+ queryResult.close();
+ } catch (RuntimeException e) {
+ throw new UDFException("Failed to close internal query result", e);
+ }
+ }
+
+ private void ensureOpen() throws UDFException {
+ if (closed) {
+ throw new UDFException("UDFResultSet is already closed");
+ }
+ }
+
+ private static List<Type> buildColumnTypes(List<ColumnHeader> columnHeaders) {
+ return columnHeaders.stream()
+ .map(ColumnHeader::getColumnType)
+ .map(UDFDataTypeTransformer::transformToUDFDataType)
+ .map(UDFDataTypeTransformer::transformUDFDataTypeToReadType)
+ .collect(Collectors.toList());
+ }
+
+ private static int[] buildColumnIndexes(List<Integer> columnIndex2TsBlockColumnIndexList) {
+ if (columnIndex2TsBlockColumnIndexList == null
+ || columnIndex2TsBlockColumnIndexList.isEmpty()) {
+ return null;
+ }
+ int size = columnIndex2TsBlockColumnIndexList.size();
+ int[] indexes = new int[size];
+ for (int i = 0; i < size; i++) {
+ indexes[i] = columnIndex2TsBlockColumnIndexList.get(i);
+ }
+ return indexes;
+ }
+
+ private List<Column> extractColumns(TsBlock tsBlock) {
+ Column[] valueColumns = tsBlock.getValueColumns();
+ if (valueColumns.length == 0) {
+ return Arrays.asList(valueColumns);
+ }
+
+ if (columnIndexes != null) {
+ return Arrays.asList(tsBlock.getColumns(columnIndexes));
+ }
+ return Arrays.asList(valueColumns);
+ }
+}