diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/CreateTableProbeFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/CreateTableProbeFunction.java new file mode 100644 index 0000000000000..90bbd4bac34e9 --- /dev/null +++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/CreateTableProbeFunction.java @@ -0,0 +1,29 @@ +/* + * 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; + +/** Probes CREATE TABLE rejection for IoTDBLocal integration tests. */ +public class CreateTableProbeFunction extends MutatingQueryProbeFunction { + + @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 0000000000000..ae2a9cc4138ab --- /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 getArgumentsSpecifications() { + return Collections.singletonList( + ScalarParameterSpecification.builder() + .name(PREFIX_PARAMETER_NAME) + .type(Type.STRING) + .defaultValue("") + .build()); + } + + @Override + public TableFunctionAnalysis analyze(Map 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 columnBuilders) { + throw new UnsupportedOperationException("framework invokes process(List, IoTDBLocal)"); + } + + @Override + public void process(List 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 0000000000000..1a036d4722a14 --- /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 getArgumentsSpecifications() { + return Collections.singletonList( + ScalarParameterSpecification.builder() + .name(PREFIX_PARAMETER_NAME) + .type(Type.STRING) + .defaultValue("") + .build()); + } + + @Override + public TableFunctionAnalysis analyze(Map 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 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 columnBuilders) { + writeDeviceIds(columnBuilders); + } + + @Override + public void process(List columnBuilders, IoTDBLocal local) { + writeDeviceIds(columnBuilders); + } + + private void writeDeviceIds(List 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 0000000000000..af49b63d2f618 --- /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 getArgumentsSpecifications() { + return Collections.singletonList( + TableParameterSpecification.builder().name(TABLE_PARAM).rowSemantics().build()); + } + + @Override + public TableFunctionAnalysis analyze(Map 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 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 idToName = Map.of(); + + @Override + public void beforeStart(IoTDBLocal local) throws UDFException { + Map 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 properColumnBuilders, + ColumnBuilder passThroughIndexBuilder) { + writeDeviceName(input, properColumnBuilders); + } + + @Override + public void process( + Record input, + List properColumnBuilders, + ColumnBuilder passThroughIndexBuilder, + IoTDBLocal local) { + writeDeviceName(input, properColumnBuilders); + } + + private void writeDeviceName(Record input, List 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 0000000000000..7839fa6b99a7d --- /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 getArgumentsSpecifications() { + return Collections.singletonList( + TableParameterSpecification.builder().name(TABLE_PARAM).rowSemantics().build()); + } + + @Override + public TableFunctionAnalysis analyze(Map 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 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 properColumnBuilders, + ColumnBuilder passThroughIndexBuilder) { + throw new UnsupportedOperationException( + "framework invokes process(Record, List, ColumnBuilder, IoTDBLocal)"); + } + + @Override + public void process( + Record input, + List 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 0000000000000..d08d0d0affec1 --- /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 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 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 0000000000000..a400f70dcaf1d --- /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 idToName = Map.of(); + private Map 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 0000000000000..c0ef353e4c1f1 --- /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 idToName = Map.of(); + private Map 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/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DropTableProbeFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DropTableProbeFunction.java new file mode 100644 index 0000000000000..a5825eab4fa8a --- /dev/null +++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/DropTableProbeFunction.java @@ -0,0 +1,29 @@ +/* + * 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; + +/** Probes DROP TABLE rejection for IoTDBLocal integration tests. */ +public class DropTableProbeFunction extends MutatingQueryProbeFunction { + + @Override + protected String mutatingSql() { + return "DROP TABLE device_info"; + } +} diff --git a/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/InsertProbeFunction.java b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/InsertProbeFunction.java new file mode 100644 index 0000000000000..a4d5fc525c4f3 --- /dev/null +++ b/integration-test/src/main/java/org/apache/iotdb/db/query/udf/example/relational/iotdblocal/InsertProbeFunction.java @@ -0,0 +1,29 @@ +/* + * 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; + +/** Probes INSERT rejection for IoTDBLocal integration tests. */ +public class InsertProbeFunction extends MutatingQueryProbeFunction { + + @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 0000000000000..7b7630bf5b046 --- /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 0000000000000..720828dcf1b39 --- /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 0000000000000..970b25c3e84a4 --- /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 0000000000000..39dfa05775185 --- /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 getArgumentsSpecifications() { + return Collections.singletonList( + ScalarParameterSpecification.builder() + .name(INPUT_PARAMETER_NAME) + .type(Type.STRING) + .build()); + } + + @Override + public TableFunctionAnalysis analyze(Map 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 columnBuilders) { + for (String value : input.split(",")) { + columnBuilders.get(0).writeBinary(new Binary(value, Charset.defaultCharset())); + } + finish = true; + } + + @Override + public void process(List 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 0000000000000..16e924c7bbd5d --- /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 0000000000000..71dd52485145c --- /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 0000000000000..1f1005925cd18 --- /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 0000000000000..7e66de05bc211 --- /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 0000000000000..bf2dbaae52628 --- /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 0000000000000..09c1784800e73 --- /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 0000000000000..de3b8a13aae8b --- /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 0000000000000..4e918e8bbef99 --- /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 0000000000000..c045357296aeb --- /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 34f57b2940b6f..6174f06bb5b3d 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 0000000000000..b84cb45965a6b --- /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/UDFResultSet.java b/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/UDFResultSet.java new file mode 100644 index 0000000000000..ed6d13442bf2b --- /dev/null +++ b/iotdb-api/udf-api/src/main/java/org/apache/iotdb/udf/api/UDFResultSet.java @@ -0,0 +1,45 @@ +/* + * 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; +import org.apache.iotdb.udf.api.relational.access.Record; + +/** Streaming result set returned by {@link IoTDBLocal#query(String)}. */ +public interface UDFResultSet extends AutoCloseable { + + /** + * @return {@code true} if another row is available + * @throws UDFException if underlying read fails + */ + boolean hasNext() throws UDFException; + + /** + * @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; + + /** 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 6c9d385ac2467..6c29de6e7ad9c 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 @@ default void beforeStart(FunctionArguments arguments) throws UDFException { // 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 @@ default void beforeStart(FunctionArguments arguments) throws UDFException { */ 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 @@ default void remove(State state, Record input) { 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 68d7793093d2a..e1daab6e4b0c1 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 @@ default void beforeStart(FunctionArguments arguments) throws UDFException { // 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 @@ default void beforeStart(FunctionArguments arguments) throws UDFException { */ 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 0fe40fc894b82..e4960fce409c0 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 @@ default void beforeStart() { // 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 @@ void process( List properColumnBuilders, ColumnBuilder passThroughIndexBuilder); + /** + * Same as {@link #process(Record, List, ColumnBuilder)} with access to {@link IoTDBLocal} for + * embedded queries. + */ + default void process( + Record input, + List 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 @@ default void finish( // do nothing } + /** + * Same as {@link #finish(List, ColumnBuilder)} with access to {@link IoTDBLocal} for embedded + * queries. + */ + default void finish( + List 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 ab8763f3c4ae7..6da0c8d5149c9 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 @@ default void beforeStart() { // 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 @@ default void beforeStart() { */ void process(List columnBuilders); + /** Same as {@link #process(List)} with access to {@link IoTDBLocal} for embedded queries. */ + default void process(List 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 @@ default void beforeStart() { 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 ced20f7bbf0b6..589e175e146e7 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.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 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 outputDataTypes) { + List 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 CommonOperatorContext getOperatorContext() { public TsBlock next() throws Exception { List 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 boolean hasNext() throws Exception { @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 e4f11d7cd49fa..d4bc8a0da4c79 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.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.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 TableFunctionOperator( List passThroughChannels, boolean isDeclaredAsPassThrough, List 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 TableFunctionOperator( this.partitionCache = new PartitionCache(); this.resultTsBlocks = new LinkedList<>(); this.requireRecordSnapshot = requireRecordSnapshot; + this.ioTDBLocal = ioTDBLocal; } @Override @@ -159,7 +166,7 @@ public TsBlock next() throws Exception { 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 TsBlock next() throws Exception { 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 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 @@ private void consumeCurrentSourceTsBlock() { 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 void close() throws Exception { 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 a37f1bef0c64d..f9d35cf67f529 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.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.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.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 static TableAccumulator createAccumulator( Set 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 inputDataTypes, + List inputExpressions, + Map inputAttributes, + boolean ascending, + boolean isAggTableScan, + String timeColumnName, + Set 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 static GroupedAccumulator createGroupedAccumulator( 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 inputDataTypes, + List inputExpressions, + Map 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 static GroupedAccumulator createGroupedAccumulator( } private static TableAccumulator createUDAFAccumulator( - String functionName, List inputDataTypes, Map inputAttributes) { + String functionName, + List inputDataTypes, + Map 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 inputDataTypes, Map inputAttributes) { + String functionName, + List inputDataTypes, + Map 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 cd6d0dc1101e9..13c0b92305f12 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.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 inputDataTypes; - private final State state; + private State state; + private final IoTDBLocal ioTDBLocal; + private boolean init; public UserDefinedAggregateFunctionAccumulator( AggregateFunctionAnalysis analysis, AggregateFunction aggregateFunction, - List inputDataTypes) { + FunctionArguments functionArguments, + List inputDataTypes, + IoTDBLocal ioTDBLocal) { + this(analysis, aggregateFunction, functionArguments, inputDataTypes, ioTDBLocal, false); + } + + private UserDefinedAggregateFunctionAccumulator( + AggregateFunctionAnalysis analysis, + AggregateFunction aggregateFunction, + FunctionArguments functionArguments, + List 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 long getEstimatedSize() { @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 void addIntermediate(Column argument) { 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 void evaluateIntermediate(ColumnBuilder columnBuilder) { @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 void addStatistics(Statistics[] statistics) { @Override public void reset() { + if (!init) { + return; + } state.reset(); } @@ -150,7 +196,10 @@ public boolean removable() { @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 2ef228b06624b..44df7fab1d7e7 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.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.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 stateArray; private final List inputDataTypes; + private final IoTDBLocal ioTDBLocal; + private boolean init = false; public GroupedUserDefinedAggregateAccumulator( - AggregateFunction aggregateFunction, List inputDataTypes) { + AggregateFunction aggregateFunction, + FunctionArguments functionArguments, + List 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 @@ private State getOrCreateState(int groupId) { @Override public void addInput(int[] groupIds, Column[] arguments, AggregationMask mask) { + initIfNeeded(); RecordIterator iterator = mask.isSelectAll() ? new RecordIterator( @@ -90,7 +117,7 @@ public void addInput(int[] groupIds, Column[] arguments, AggregationMask mask) { 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 void addInput(int[] groupIds, Column[] arguments, AggregationMask mask) { 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 void addIntermediate(int[] groupIds, Column argument) { 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 void evaluateIntermediate(int groupId, ColumnBuilder columnBuilder) { @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 void reset() { @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 a25188361af2d..7d3c1ab9bc547 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 @@ private ColumnTransformer getFunctionColumnTransformer( .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 static class Context { @SuppressWarnings("unused") private final Optional memoryReservationManager; + private final String fragmentInstanceId; + + private final String outerGlobalQueryId; + + private final long outerQueryDeadlineMs; + + @Nullable private final IoTDBLocalFactory ioTDBLocalFactory; + public Context( SessionInfo sessionInfo, List leafList, @@ -1966,6 +1974,40 @@ public Context( 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 leafList, + Map> inputLocations, + Map cache, + Map hasSeen, + List commonTransformerList, + List 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 Context( 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 bfac196461e3d..86786404bc4c3 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.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.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 @@ protected Operator constructFilterAndProjectOperator( PlanNodeId planNodeId, C context) { + String fragmentInstanceId = getFragmentInstanceId(context); + String outerGlobalQueryId = getQueryId(context); + long outerQueryDeadlineMs = getOuterQueryDeadlineMs(context); + IoTDBLocalFactory ioTDBLocalFactory = getIoTDBLocalFactory(context); + final List filterOutputDataTypes = new ArrayList<>(inputDataTypes); // records LeafColumnTransformer of filter @@ -341,7 +348,11 @@ protected Operator constructFilterAndProjectOperator( 0, context.getTableTypeProvider(), metadata, - context.getMemoryReservationManager()); + context.getMemoryReservationManager(), + fragmentInstanceId, + outerGlobalQueryId, + outerQueryDeadlineMs, + ioTDBLocalFactory); return visitor.process(p, filterColumnTransformerContext); }) @@ -369,7 +380,11 @@ protected Operator constructFilterAndProjectOperator( inputLocations.size(), context.getTableTypeProvider(), metadata, - context.getMemoryReservationManager()); + context.getMemoryReservationManager(), + fragmentInstanceId, + outerGlobalQueryId, + outerQueryDeadlineMs, + ioTDBLocalFactory); for (Expression expression : projectExpressions) { projectOutputTransformerList.add( @@ -393,6 +408,22 @@ protected Operator constructFilterAndProjectOperator( 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 @@ protected AggregationOperator planGlobalAggregation( false, null, Collections.emptySet(), - operatorContext.getMemoryReservationContext()))); + operatorContext.getMemoryReservationContext(), + context))); return createAggregationOperator(operatorContext, child, aggregatorBuilder.build()); } @@ -1344,7 +1376,8 @@ protected TableAggregator buildAggregator( boolean isAggTableScan, String timeColumnName, Set measurementColumnNames, - MemoryReservationManager memoryReservationManager) { + MemoryReservationManager memoryReservationManager, + C context) { List argumentChannels = new ArrayList<>(); for (Expression argument : aggregation.getArguments()) { Symbol argumentSymbol = Symbol.from(argument); @@ -1356,6 +1389,10 @@ protected TableAggregator buildAggregator( 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 @@ protected TableAggregator buildAggregator( timeColumnName, measurementColumnNames, aggregation.isDistinct(), - memoryReservationManager); + memoryReservationManager, + ioTDBLocal); OptionalInt maskChannel = OptionalInt.empty(); if (aggregation.hasMask()) { @@ -1411,7 +1449,8 @@ protected Operator planGroupByAggregation( false, null, Collections.emptySet(), - context.getMemoryReservationManager()))); + context.getMemoryReservationManager(), + context))); CommonOperatorContext operatorContext = addOperatorContext( @@ -1439,7 +1478,8 @@ protected Operator planGroupByAggregation( v, node.getStep(), typeProvider, - context.getMemoryReservationManager()))); + context.getMemoryReservationManager(), + context))); Set preGroupedKeys = ImmutableSet.copyOf(node.getPreGroupedSymbols()); List groupingKeys = node.getGroupingKeys(); @@ -1496,7 +1536,8 @@ protected Operator planGroupByAggregation( v, node.getStep(), typeProvider, - context.getMemoryReservationManager()))); + context.getMemoryReservationManager(), + context))); CommonOperatorContext operatorContext = addOperatorContext( context, node.getPlanNodeId(), HashAggregationOperator.class.getSimpleName()); @@ -1621,7 +1662,8 @@ protected GroupedAggregator buildGroupByAggregator( AggregationNode.Aggregation aggregation, AggregationNode.Step step, ITableTypeProvider typeProvider, - MemoryReservationManager memoryReservationManager) { + MemoryReservationManager memoryReservationManager, + C context) { List argumentChannels = new ArrayList<>(); for (Expression argument : aggregation.getArguments()) { Symbol argumentSymbol = Symbol.from(argument); @@ -1633,6 +1675,10 @@ protected GroupedAggregator buildGroupByAggregator( 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 @@ protected GroupedAggregator buildGroupByAggregator( Collections.emptyMap(), true, aggregation.isDistinct(), - memoryReservationManager); + memoryReservationManager, + ioTDBLocal); OptionalInt maskChannel = OptionalInt.empty(); if (aggregation.hasMask()) { @@ -1671,7 +1718,9 @@ public Operator visitTableFunctionProcessor(TableFunctionProcessorNode node, C c 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 Operator visitTableFunctionProcessor(TableFunctionProcessorNode node, C c } else { partitionChannels = Collections.emptyList(); } + IoTDBLocal ioTDBLocal = createIoTDBLocal(context); return new TableFunctionOperator( operatorContext, processorProvider, @@ -1729,7 +1779,8 @@ public Operator visitTableFunctionProcessor(TableFunctionProcessorNode node, C c .map(TableFunctionNode.PassThroughSpecification::isDeclaredAsPassThrough) .orElse(false), partitionChannels, - node.isRequireRecordSnapshot()); + node.isRequireRecordSnapshot(), + ioTDBLocal); } } @@ -2467,4 +2518,40 @@ public Operator visitTopKRanking(TopKRankingNode node, C context) { } 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 4e22315b62868..3288dcbad0768 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 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 inputTypes; + private final IoTDBLocal ioTDBLocal; + private boolean init = false; public UserDefineScalarFunctionTransformer( Type returnType, ScalarFunction scalarFunction, - List childrenTransformers) { + List 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 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 @@ protected void doTransform( @Override protected void doTransform( List 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 @@ protected void doTransform( builder.appendNull(); continue; } - Object result = scalarFunction.evaluate(input); + Object result = scalarFunction.evaluate(input, ioTDBLocal); if (result == null) { builder.appendNull(); } else { @@ -97,8 +132,11 @@ protected void doTransform( @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 b3fa3c049f363..21dbffb606ebd 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 IClientSession getCurrSession() { 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 54ac31ce7f774..d940d2260968f 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 @@ private TSExecuteStatementResp executeStatementInternal( } } - private void clearUp( + public static void clearUp( IClientSession clientSession, Long statementId, Long queryId, @@ -572,7 +572,7 @@ private void clearUp( 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 6bda6e9c14d6e..595e9da7065fb 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 @@ -152,6 +152,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 dataNodeQueryContextMap; private DataNodeQueryContext dataNodeQueryContext; @@ -206,6 +209,7 @@ public static FragmentInstanceContext createFragmentInstanceContext( IDataRegionForQuery dataRegion, TimePredicate globalTimePredicate, Map dataNodeQueryContextMap, + long outerQueryDeadlineMs, boolean debug, boolean isVerbose) { FragmentInstanceContext instanceContext = @@ -216,6 +220,7 @@ public static FragmentInstanceContext createFragmentInstanceContext( dataRegion, globalTimePredicate, dataNodeQueryContextMap, + outerQueryDeadlineMs, debug, isVerbose); instanceContext.initialize(); @@ -271,6 +276,7 @@ private FragmentInstanceContext( IDataRegionForQuery dataRegion, TimePredicate globalTimePredicate, Map dataNodeQueryContextMap, + long outerQueryDeadlineMs, boolean debug, boolean verbose) { super(debug, verbose); @@ -278,6 +284,7 @@ private FragmentInstanceContext( this.stateMachine = stateMachine; this.executionEndTime.set(END_TIME_INITIAL_VALUE); this.sessionInfo = sessionInfo; + this.outerQueryDeadlineMs = outerQueryDeadlineMs; this.dataRegion = dataRegion; this.globalTimeFilter = globalTimePredicate == null @@ -567,6 +574,10 @@ public SessionInfo getSessionInfo() { return sessionInfo; } + public long getOuterQueryDeadlineMs() { + return outerQueryDeadlineMs; + } + public Optional 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 e5f3cbc7e0d58..a9473982237e6 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.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 FragmentInstanceInfo execDataQueryFragmentInstance( 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 c85d00e17e16e..285c67dd7babc 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 DriverTask( 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 DriverTask( 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 73a09597cd594..1e89784aa0add 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 @@ -155,6 +155,7 @@ 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; @@ -307,6 +308,17 @@ private ExecutionResult execution( boolean userQuery, boolean debug, BiFunction 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 iQueryExecutionFactory) { long startTime = System.currentTimeMillis(); QueryId globalQueryId = queryIdGenerator.createNextQueryId(); MPPQueryContext queryContext = null; @@ -325,6 +337,9 @@ private ExecutionResult execution( 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 { @@ -488,12 +503,39 @@ public ExecutionResult executeForTableModel( 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 185c0d12c0632..7d628e78d7933 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.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; @@ -132,6 +133,7 @@ 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; @@ -1464,7 +1466,8 @@ public Operator visitNonAlignedAggregationTreeDeviceViewScan( true, timeColumnName, measurementColumnsIndexMap.keySet(), - context.getMemoryReservationManager())); + context.getMemoryReservationManager(), + context)); } ITableTimeRangeIterator timeRangeIterator = null; @@ -2101,4 +2104,24 @@ private OptimizeType canUseLastCacheOptimize( 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 6590870839472..4d78236f38d1b 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 FragmentInstanceId getFragmentInstanceId() { return driverContext.getFragmentInstanceContext().getId(); } + public long getOuterQueryDeadlineMs() { + return driverContext.getFragmentInstanceContext().getOuterQueryDeadlineMs(); + } + public List 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 0000000000000..5abcf05d22b3f --- /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 0000000000000..419ea00eb0028 --- /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 0000000000000..630360f485808 --- /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 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 0000000000000..74cc1a3859f77 --- /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 openResultSets; + private final int index; + private final InternalQueryResult queryResult; + private final List columnTypes; + + /** null when output columns align with TsBlock value columns without index mapping */ + private final int[] columnIndexes; + + private Iterator rowIterator; + private boolean closed; + + public UDFResultSetImpl( + List 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 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 buildColumnTypes(List columnHeaders) { + return columnHeaders.stream() + .map(ColumnHeader::getColumnType) + .map(UDFDataTypeTransformer::transformToUDFDataType) + .map(UDFDataTypeTransformer::transformUDFDataTypeToReadType) + .collect(Collectors.toList()); + } + + private static int[] buildColumnIndexes(List 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 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); + } +}