module-3 Signed-off-by: Weihao Li <18110526956@163.com>
diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index 3feefd7..5cb613d 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java
@@ -3818,5 +3818,8 @@ public static final String EXCEPTION_VISIBLEALIASES_IS_NULL_630B27F1 = "visibleAliases is null"; public static final String EXCEPTION_HAS_NO_PERMISSION_TO_EXECUTE_ARG_BECAUSE_ONLY_THE_SUPERUSER_CAN_ALTER_HIM_HERSELF_C5902893 = "Has no permission to execute %s, because only the superuser can alter him/herself."; + public static final String + LOG_FAILED_TO_CLEAN_DEVICEENTRY_DATA_SET_ASYNCHRONOUSLY_QUERYID_ARG_PLANNODEID_ARG_9106C4C5 = + "Failed to clean DeviceEntry data set asynchronously: queryId=%s, planNodeId=%s"; }
diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index 0d4e7cc..89bafbd 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java
@@ -4575,5 +4575,8 @@ public static final String EXCEPTION_VISIBLEALIASES_IS_NULL_630B27F1 = "visibleAliases 不能为空"; public static final String EXCEPTION_HAS_NO_PERMISSION_TO_EXECUTE_ARG_BECAUSE_ONLY_THE_SUPERUSER_CAN_ALTER_HIM_HERSELF_C5902893 = "无权执行 %s，因为只有超级用户可以修改其自身。"; + public static final String + LOG_FAILED_TO_CLEAN_DEVICEENTRY_DATA_SET_ASYNCHRONOUSLY_QUERYID_ARG_PLANNODEID_ARG_9106C4C5 = + "异步清理 DeviceEntry 数据集失败：queryId=%s，planNodeId=%s"; }
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java index 2d3447d..0e0d6b1 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java
@@ -43,11 +43,14 @@ import org.apache.iotdb.db.queryengine.execution.memory.LocalMemoryManager; import org.apache.iotdb.db.queryengine.metric.DataExchangeCostMetricSet; import org.apache.iotdb.db.queryengine.metric.DataExchangeCountMetricSet; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.DeviceEntrySpillManager; import org.apache.iotdb.db.utils.SetThreadName; import org.apache.iotdb.mpp.rpc.thrift.MPPDataExchangeService; import org.apache.iotdb.mpp.rpc.thrift.TAcknowledgeDataBlockEvent; import org.apache.iotdb.mpp.rpc.thrift.TCloseSinkChannelEvent; import org.apache.iotdb.mpp.rpc.thrift.TEndOfDataBlockEvent; +import org.apache.iotdb.mpp.rpc.thrift.TFetchDeviceEntrySegmentReq; +import org.apache.iotdb.mpp.rpc.thrift.TFetchDeviceEntrySegmentResp; import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId; import org.apache.iotdb.mpp.rpc.thrift.TGetDataBlockRequest; import org.apache.iotdb.mpp.rpc.thrift.TGetDataBlockResponse; @@ -97,6 +100,46 @@ DataExchangeCountMetricSet.getInstance(); @Override + public TFetchDeviceEntrySegmentResp fetchDeviceEntrySegment( + TFetchDeviceEntrySegmentReq request) { + try { + DeviceEntrySpillManager spillManager = DeviceEntrySpillManager.getInstance(); + byte[] payload = + spillManager.readSegment( + request.getQueryId(), request.getPlanNodeId(), request.getSegmentId()); + if (request.getSegmentId() > 0) { + spillManager.deleteSegment( + request.getQueryId(), request.getPlanNodeId(), request.getSegmentId() - 1); + } + return new TFetchDeviceEntrySegmentResp( + new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode())) + .setPayload(payload); + } catch (IOException | RuntimeException e) { + return new TFetchDeviceEntrySegmentResp( + new TSStatus(TSStatusCode.INTERNAL_SERVER_ERROR.getStatusCode()) + .setMessage(e.getMessage())); + } + } + + @Override + public TSStatus finishDeviceEntrySegment(String queryId, String planNodeId) { + executorService.submit( + () -> { + try { + DeviceEntrySpillManager.getInstance().finishSegmentDataSet(queryId, planNodeId); + } catch (IOException | RuntimeException e) { + LOGGER.warn( + DataNodeQueryMessages + .LOG_FAILED_TO_CLEAN_DEVICEENTRY_DATA_SET_ASYNCHRONOUSLY_QUERYID_ARG_PLANNODEID_ARG_9106C4C5, + queryId, + planNodeId, + e); + } + }); + return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()); + } + + @Override public TGetDataBlockResponse getDataBlock(TGetDataBlockRequest req) throws TException { long startTime = System.nanoTime(); try (SetThreadName fragmentInstanceName = @@ -623,6 +666,11 @@ return mppDataExchangeService; } + public IClientManager<TEndPoint, SyncDataNodeMPPDataExchangeServiceClient> + getMppDataExchangeServiceClientManager() { + return mppDataExchangeServiceClientManager; + } + public void deRegisterFragmentInstanceFromMemoryPool( String queryId, String fragmentInstanceId, boolean forceDeregister) { localMemoryManager
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/BatchDeviceEntrySource.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/BatchDeviceEntrySource.java new file mode 100644 index 0000000..72f00e7 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/BatchDeviceEntrySource.java
@@ -0,0 +1,35 @@ +/* + * 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.plan.relational.metadata.spill; + +import org.apache.iotdb.db.queryengine.plan.relational.metadata.DeviceEntry; + +import java.io.IOException; +import java.util.List; + +public interface BatchDeviceEntrySource extends AutoCloseable { + + boolean hasNextBatch(); + + List<DeviceEntry> nextBatch() throws IOException; + + @Override + void close() throws IOException; +}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryDataSetHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryDataSetHandle.java index d731aae..f295cd0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryDataSetHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryDataSetHandle.java
@@ -101,7 +101,7 @@ new PlanNodeId(ReadWriteIOUtils.readString(byteBuffer)), ThriftCommonsSerDeUtils.deserializeTEndPoint(byteBuffer), ReadWriteIOUtils.readInt(byteBuffer), - ReadWriteIOUtils.readLong(byteBuffer), + ReadWriteIOUtils.readInt(byteBuffer), ReadWriteIOUtils.readBool(byteBuffer)); } }
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryReader.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryReader.java index 5fa2e04..cf5410d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryReader.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryReader.java
@@ -25,10 +25,10 @@ public interface DeviceEntryReader extends AutoCloseable { - public boolean hasNext() throws IOException; + boolean hasNext() throws IOException; - public DeviceEntry next() throws IOException; + DeviceEntry next() throws IOException; @Override - public void close() throws IOException; + void close() throws IOException; }
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryRpcSegmentFetcher.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryRpcSegmentFetcher.java new file mode 100644 index 0000000..90b33b5 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryRpcSegmentFetcher.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.queryengine.plan.relational.metadata.spill; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.client.IClientManager; +import org.apache.iotdb.commons.client.exception.ClientManagerException; +import org.apache.iotdb.commons.client.sync.SyncDataNodeMPPDataExchangeServiceClient; +import org.apache.iotdb.commons.utils.TestOnly; +import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeService; +import org.apache.iotdb.mpp.rpc.thrift.TFetchDeviceEntrySegmentReq; +import org.apache.iotdb.mpp.rpc.thrift.TFetchDeviceEntrySegmentResp; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.apache.thrift.TException; + +import java.io.IOException; + +public final class DeviceEntryRpcSegmentFetcher implements DeviceEntrySegmentFetcher { + + private static final int MAX_ATTEMPTS = 3; + + private final IClientManager<TEndPoint, SyncDataNodeMPPDataExchangeServiceClient> clientManager; + + private DeviceEntryRpcSegmentFetcher() { + this( + MPPDataExchangeService.getInstance() + .getMPPDataExchangeManager() + .getMppDataExchangeServiceClientManager()); + } + + @TestOnly + public DeviceEntryRpcSegmentFetcher( + IClientManager<TEndPoint, SyncDataNodeMPPDataExchangeServiceClient> clientManager) { + this.clientManager = clientManager; + } + + public static DeviceEntryRpcSegmentFetcher getInstance() { + return DeviceEntryRpcSegmentFetcherHolder.INSTANCE; + } + + @Override + public byte[] fetch(DeviceEntryDataSetHandle handle, int segmentId) throws IOException { + IOException failure = null; + for (int attempt = 0; attempt < MAX_ATTEMPTS; attempt++) { + TFetchDeviceEntrySegmentResp response; + try { + try (SyncDataNodeMPPDataExchangeServiceClient client = + clientManager.borrowClient(handle.getCoordinatorEndPoint())) { + response = client.fetchDeviceEntrySegment(createFetchRequest(handle, segmentId)); + } + } catch (ClientManagerException | TException e) { + failure = new IOException(e); + continue; + } + if (response.getStatus().getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + throw new IOException(response.getStatus().getMessage()); + } + return response.getPayload(); + } + throw failure; + } + + @Override + public void finish(DeviceEntryDataSetHandle handle) { + for (int attempt = 0; attempt < MAX_ATTEMPTS; attempt++) { + try { + try (SyncDataNodeMPPDataExchangeServiceClient client = + clientManager.borrowClient(handle.getCoordinatorEndPoint())) { + org.apache.iotdb.common.rpc.thrift.TSStatus status = + client.finishDeviceEntrySegment(handle.getQueryId(), handle.getPlanNodeId().getId()); + if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + return; + } + } + return; + } catch (ClientManagerException | TException e) { + // Cleanup notification is best effort and must not fail the query. + } + } + } + + private TFetchDeviceEntrySegmentReq createFetchRequest( + DeviceEntryDataSetHandle handle, int segmentId) { + return new TFetchDeviceEntrySegmentReq( + handle.getQueryId(), handle.getPlanNodeId().getId(), segmentId); + } + + private static class DeviceEntryRpcSegmentFetcherHolder { + private static final DeviceEntryRpcSegmentFetcher INSTANCE = new DeviceEntryRpcSegmentFetcher(); + + private DeviceEntryRpcSegmentFetcherHolder() {} + } +}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntrySegmentFetcher.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntrySegmentFetcher.java new file mode 100644 index 0000000..c66e46f --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntrySegmentFetcher.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.queryengine.plan.relational.metadata.spill; + +import java.io.IOException; + +public interface DeviceEntrySegmentFetcher { + + byte[] fetch(DeviceEntryDataSetHandle handle, int segmentId) throws IOException; + + void finish(DeviceEntryDataSetHandle handle); +}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntrySpillManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntrySpillManager.java index 5580c08..2b28c95 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntrySpillManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntrySpillManager.java
@@ -45,7 +45,7 @@ } public Path register(String queryId, PlanNodeId planNodeId) throws IOException { - Path ownerDirectory = rootDirectory().resolve(queryId).resolve(planNodeId.getId()); + Path ownerDirectory = resolveOwnerDirectory(queryId, planNodeId.getId()); Files.createDirectories(ownerDirectory); queryDirectories .computeIfAbsent(queryId, ignored -> ConcurrentHashMap.newKeySet()) @@ -66,7 +66,7 @@ public void deregisterQuery(String queryId) throws IOException { queryDirectories.remove(queryId); - FileUtils.deleteDirectory(rootDirectory().resolve(queryId).toFile()); + FileUtils.deleteDirectory(resolveQueryDirectory(queryId).toFile()); } @TestOnly @@ -109,17 +109,8 @@ } public void finishSegmentDataSet(String queryId, String planNodeId) throws IOException { - deregisterOwner(queryId, rootDirectory().resolve(queryId).resolve(planNodeId)); - } - - public void deregisterFragment(String queryId, String fragmentInstanceId) throws IOException { - FileUtils.deleteDirectory( - resolveUnderRoot(fragmentRootDirectory(), queryId, fragmentInstanceId).toFile()); - } - - public void clearStaleFragmentData() throws IOException { - FileUtils.deleteDirectory(fragmentRootDirectory().toFile()); - Files.createDirectories(fragmentRootDirectory()); + Path ownerDirectory = resolveOwnerDirectory(queryId, planNodeId); + deregisterOwner(queryId, ownerDirectory); } private Path resolveRegisteredDataSetDirectory(String queryId, String dataSetId) @@ -130,7 +121,7 @@ .anyMatch(path -> path.toString().equals("..") || path.toString().equals("."))) { throw new IllegalArgumentException(); } - Path queryDirectory = rootDirectory().resolve(queryId).normalize(); + Path queryDirectory = resolveQueryDirectory(queryId); Path dataSetDirectory = queryDirectory.resolve(relativeDataSetPath).resolve("fi").normalize(); if (!dataSetDirectory.startsWith(queryDirectory)) { throw new IllegalArgumentException(); @@ -166,20 +157,16 @@ return Path.of(IoTDBDescriptor.getInstance().getConfig().getSortTmpDir(), "device-entry"); } - private Path fragmentRootDirectory() { - return rootDirectory().resolve("fragment"); + private Path resolveOwnerDirectory(String queryId, String planNodeId) { + Path queryDirectory = resolveQueryDirectory(queryId); + Path ownerDirectory = queryDirectory.resolve(planNodeId).normalize(); + return ownerDirectory; } - private Path resolveUnderRoot(Path root, String... children) { - Path result = root; - for (String child : children) { - result = result.resolve(child); - } - result = result.normalize(); - if (!result.startsWith(root.normalize())) { - throw new IllegalArgumentException(); - } - return result; + private Path resolveQueryDirectory(String queryId) { + Path root = rootDirectory().normalize(); + Path queryDirectory = root.resolve(queryId).normalize(); + return queryDirectory; } private static class DeviceEntrySpillManagerHolder {
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/InMemoryDeviceEntrySource.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/InMemoryDeviceEntrySource.java new file mode 100644 index 0000000..8d67e76 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/InMemoryDeviceEntrySource.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.queryengine.plan.relational.metadata.spill; + +import org.apache.iotdb.db.queryengine.plan.relational.metadata.DeviceEntry; + +import java.util.Collections; +import java.util.List; + +public final class InMemoryDeviceEntrySource implements BatchDeviceEntrySource { + + private List<DeviceEntry> entries; + + public InMemoryDeviceEntrySource(List<DeviceEntry> entries) { + this.entries = entries; + } + + @Override + public boolean hasNextBatch() { + return entries != null; + } + + @Override + public List<DeviceEntry> nextBatch() { + if (entries == null) { + return Collections.emptyList(); + } + List<DeviceEntry> result = entries; + entries = null; + return result; + } + + @Override + public void close() { + entries = null; + } +}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/LocalSegmentDeviceEntrySource.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/LocalSegmentDeviceEntrySource.java new file mode 100644 index 0000000..684b880 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/LocalSegmentDeviceEntrySource.java
@@ -0,0 +1,64 @@ +/* + * 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.plan.relational.metadata.spill; + +import org.apache.iotdb.db.queryengine.plan.relational.metadata.DeviceEntry; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; + +public final class LocalSegmentDeviceEntrySource extends SegmentDeviceEntrySource { + + private final DeviceEntrySpillManager spillManager; + + public LocalSegmentDeviceEntrySource(DeviceEntryDataSetHandle handle) { + this(handle, DeviceEntrySpillManager.getInstance()); + } + + public LocalSegmentDeviceEntrySource( + DeviceEntryDataSetHandle handle, DeviceEntrySpillManager spillManager) { + super(handle); + this.spillManager = spillManager; + } + + @Override + public List<DeviceEntry> nextBatch() throws IOException { + int segmentId = nextSegmentId; + Path segment = acquireSegment(segmentId); + List<DeviceEntry> result = deserialize(Files.readAllBytes(segment)); + releaseSegment(segmentId); + nextSegmentId++; + return result; + } + + private Path acquireSegment(int segmentId) throws IOException { + return spillManager.resolveSegment(handle.getQueryId(), handle.getPlanNodeId(), segmentId); + } + + private void releaseSegment(int segmentId) throws IOException { + if (segmentId + 1 == handle.getSegmentCount()) { + spillManager.finishSegmentDataSet(handle.getQueryId(), handle.getPlanNodeId().getId()); + } else { + spillManager.deleteSegment(handle.getQueryId(), handle.getPlanNodeId(), segmentId); + } + } +}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/RemoteSegmentDeviceEntrySource.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/RemoteSegmentDeviceEntrySource.java new file mode 100644 index 0000000..d43d7af --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/RemoteSegmentDeviceEntrySource.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.plan.relational.metadata.spill; + +import org.apache.iotdb.db.queryengine.plan.relational.metadata.DeviceEntry; + +import java.io.IOException; +import java.util.List; + +public final class RemoteSegmentDeviceEntrySource extends SegmentDeviceEntrySource { + + private final DeviceEntrySegmentFetcher fetcher; + private boolean finished; + + public RemoteSegmentDeviceEntrySource(DeviceEntryDataSetHandle handle) { + this(handle, DeviceEntryRpcSegmentFetcher.getInstance()); + } + + public RemoteSegmentDeviceEntrySource( + DeviceEntryDataSetHandle handle, DeviceEntrySegmentFetcher fetcher) { + super(handle); + this.fetcher = fetcher; + } + + @Override + public List<DeviceEntry> nextBatch() throws IOException { + int segmentId = nextSegmentId; + byte[] payload = fetcher.fetch(handle, segmentId); + List<DeviceEntry> result = deserialize(payload); + nextSegmentId++; + if (!hasNextBatch()) { + finish(); + } + return result; + } + + private void finish() { + fetcher.finish(handle); + finished = true; + } + + @Override + public void close() throws IOException { + if (!finished) { + finish(); + } + } +}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/SegmentDeviceEntrySource.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/SegmentDeviceEntrySource.java new file mode 100644 index 0000000..0a00e95 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/SegmentDeviceEntrySource.java
@@ -0,0 +1,78 @@ +/* + * 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.plan.relational.metadata.spill; + +import org.apache.iotdb.db.queryengine.common.DataNodeEndPoints; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.DeviceEntry; + +import java.io.ByteArrayInputStream; +import java.io.DataInputStream; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +public abstract class SegmentDeviceEntrySource implements BatchDeviceEntrySource { + + protected final DeviceEntryDataSetHandle handle; + protected int nextSegmentId; + + protected SegmentDeviceEntrySource(DeviceEntryDataSetHandle handle) { + this.handle = handle; + } + + public static SegmentDeviceEntrySource create(DeviceEntryDataSetHandle handle) { + return handle.getCoordinatorEndPoint().equals(DataNodeEndPoints.LOCAL_HOST_INTERNAL_ENDPOINT) + ? new LocalSegmentDeviceEntrySource(handle) + : new RemoteSegmentDeviceEntrySource(handle); + } + + @Override + public final boolean hasNextBatch() { + return nextSegmentId < handle.getSegmentCount(); + } + + protected final List<DeviceEntry> deserialize(byte[] segmentBytes) throws IOException { + List<DeviceEntry> result = new ArrayList<>(); + int segmentLength = segmentBytes.length; + try (DataInputStream input = new DataInputStream(new ByteArrayInputStream(segmentBytes))) { + int consumedBytes = 0; + while (consumedBytes < segmentLength) { + if (segmentLength - consumedBytes < Integer.BYTES) { + throw new IOException(); + } + int length = input.readInt(); + consumedBytes += Integer.BYTES; + if (length < 0 || length > segmentLength - consumedBytes) { + throw new IOException(); + } + byte[] bytes = new byte[length]; + input.readFully(bytes); + consumedBytes += length; + result.add(DeviceEntry.deserialize(bytes)); + } + } + return result; + } + + @Override + public void close() throws IOException { + // No local cache is created by a segment source. + } +}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/QueryCardinalityUtil.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/QueryCardinalityUtil.java index a59b2e0..6b49aeb 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/QueryCardinalityUtil.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/QueryCardinalityUtil.java
@@ -218,10 +218,10 @@ && !node.getProjection().getMap().isEmpty()) { // also exist date_bin return Range.atLeast(0L); } else { - return Range.atMost(node.getDeviceEntryCount()); + return Range.atMost((long) node.getDeviceEntryCount()); } } else { - return Range.singleton(node.getDeviceEntryCount()); + return Range.singleton((long) node.getDeviceEntryCount()); } }
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryRpcSegmentFetcherTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryRpcSegmentFetcherTest.java new file mode 100644 index 0000000..8572f00 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryRpcSegmentFetcherTest.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.queryengine.plan.relational.metadata.spill; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.client.IClientManager; +import org.apache.iotdb.commons.client.sync.SyncDataNodeMPPDataExchangeServiceClient; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.mpp.rpc.thrift.TFetchDeviceEntrySegmentResp; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.apache.thrift.TException; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mockito; + +import static org.junit.Assert.assertArrayEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class DeviceEntryRpcSegmentFetcherTest { + + private IClientManager<TEndPoint, SyncDataNodeMPPDataExchangeServiceClient> clientManager; + private SyncDataNodeMPPDataExchangeServiceClient client; + private DeviceEntryRpcSegmentFetcher fetcher; + private DeviceEntryDataSetHandle handle; + + @Before + @SuppressWarnings("unchecked") + public void setUp() throws Exception { + clientManager = Mockito.mock(IClientManager.class); + client = Mockito.mock(SyncDataNodeMPPDataExchangeServiceClient.class); + when(clientManager.borrowClient(any())).thenReturn(client); + fetcher = new DeviceEntryRpcSegmentFetcher(clientManager); + handle = + new DeviceEntryDataSetHandle( + "query", new PlanNodeId("scan"), new TEndPoint("127.0.0.1", 10740), 1, 1, false); + } + + @Test + public void testFetchRetriesNetworkFailure() throws Exception { + byte[] payload = new byte[] {1, 2, 3}; + when(client.fetchDeviceEntrySegment(any())) + .thenThrow(new TException()) + .thenThrow(new TException()) + .thenReturn( + new TFetchDeviceEntrySegmentResp( + new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode())) + .setPayload(payload)); + + assertArrayEquals(payload, fetcher.fetch(handle, 0)); + verify(client, times(3)).fetchDeviceEntrySegment(any()); + } + + @Test + public void testFinishRetriesNetworkFailureAtMostThreeTimes() throws Exception { + when(client.finishDeviceEntrySegment(any(), any())).thenThrow(new TException()); + + fetcher.finish(handle); + + verify(client, times(3)).finishDeviceEntrySegment(any(), any()); + } + + @Test + public void testFinishDoesNotRetryServerFailure() throws Exception { + when(client.finishDeviceEntrySegment(any(), any())) + .thenReturn(new TSStatus(TSStatusCode.INTERNAL_SERVER_ERROR.getStatusCode())); + + fetcher.finish(handle); + + verify(client, times(1)).finishDeviceEntrySegment(any(), any()); + } +}
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/SegmentDeviceEntrySourceTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/SegmentDeviceEntrySourceTest.java new file mode 100644 index 0000000..2eba72d --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/SegmentDeviceEntrySourceTest.java
@@ -0,0 +1,176 @@ +/* + * 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.plan.relational.metadata.spill; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.AlignedDeviceEntry; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.DeviceEntry; + +import org.apache.tsfile.file.metadata.IDeviceID; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import java.io.ByteArrayOutputStream; +import java.io.DataOutputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +public class SegmentDeviceEntrySourceTest { + + private Path queryDirectory; + private String originalSortTmpDir; + + @Before + public void setUp() throws Exception { + queryDirectory = Files.createTempDirectory("device-entry-source-test"); + originalSortTmpDir = IoTDBDescriptor.getInstance().getConfig().getSortTmpDir(); + IoTDBDescriptor.getInstance().getConfig().setSortTmpDir(queryDirectory.toString()); + } + + @After + public void tearDown() throws Exception { + DeviceEntrySpillManager.getInstance().clearStaleData(); + Files.deleteIfExists(queryDirectory.resolve("device-entry")); + Files.deleteIfExists(queryDirectory); + IoTDBDescriptor.getInstance().getConfig().setSortTmpDir(originalSortTmpDir); + } + + @Test + public void testLocalSourceConsumesSegmentsAndCleansDataSet() throws Exception { + List<DeviceEntry> expected = createEntries(20); + PlanNodeId planNodeId = new PlanNodeId("scan-local"); + SpilledDeviceEntryDataSet dataSet; + try (DeviceEntryMaterializer materializer = + new DeviceEntryMaterializer("q-local", planNodeId, 128, false)) { + for (DeviceEntry entry : expected) { + materializer.append(entry); + } + materializer.forceSpill(); + dataSet = (SpilledDeviceEntryDataSet) materializer.finish(); + } + + DeviceEntryDataSetHandle handle = + new DeviceEntryDataSetHandle( + "q-local", + planNodeId, + new TEndPoint("127.0.0.1", 1), + dataSet.getSegments().size(), + expected.size(), + false); + List<DeviceEntry> actual = new ArrayList<>(); + try (LocalSegmentDeviceEntrySource source = new LocalSegmentDeviceEntrySource(handle)) { + while (source.hasNextBatch()) { + actual.addAll(source.nextBatch()); + } + } + + assertEquals(expected, actual); + assertFalse(Files.exists(queryDirectory.resolve("device-entry/q-local/scan-local"))); + } + + @Test + public void testRemoteSourceFetchesSegmentsAndFinishes() throws Exception { + List<DeviceEntry> expected = createEntries(3); + RecordingFetcher fetcher = new RecordingFetcher(expected); + DeviceEntryDataSetHandle handle = + new DeviceEntryDataSetHandle( + "q-remote", + new PlanNodeId("scan-remote"), + new TEndPoint("127.0.0.2", 2), + expected.size(), + expected.size(), + false); + List<DeviceEntry> actual = new ArrayList<>(); + try (RemoteSegmentDeviceEntrySource source = + new RemoteSegmentDeviceEntrySource(handle, fetcher)) { + while (source.hasNextBatch()) { + actual.addAll(source.nextBatch()); + } + } + + assertEquals(expected, actual); + assertEquals(List.of(0, 1, 2), fetcher.segmentIds); + assertTrue(fetcher.finished); + } + + @Test + public void testFinishUnregisteredDataSetIsIdempotent() throws Exception { + DeviceEntrySpillManager.getInstance() + .finishSegmentDataSet("unregistered-query", "unregistered-scan"); + DeviceEntrySpillManager.getInstance() + .finishSegmentDataSet("unregistered-query", "unregistered-scan"); + } + + private static List<DeviceEntry> createEntries(int count) { + List<DeviceEntry> entries = new ArrayList<>(); + for (int i = 0; i < count; i++) { + entries.add( + new AlignedDeviceEntry( + IDeviceID.Factory.DEFAULT_FACTORY.create(new String[] {"table", "device" + i}), + new org.apache.tsfile.utils.Binary[0])); + } + return entries; + } + + private static byte[] serializeSegment(DeviceEntry entry) throws Exception { + byte[] payload = entry.serializeToBytes(); + try (ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + DataOutputStream output = new DataOutputStream(bytes)) { + output.writeInt(payload.length); + output.write(payload); + return bytes.toByteArray(); + } + } + + private static final class RecordingFetcher implements DeviceEntrySegmentFetcher { + + private final List<DeviceEntry> entries; + private final List<Integer> segmentIds = new ArrayList<>(); + private boolean finished; + + private RecordingFetcher(List<DeviceEntry> entries) { + this.entries = entries; + } + + @Override + public byte[] fetch(DeviceEntryDataSetHandle handle, int segmentId) throws java.io.IOException { + segmentIds.add(segmentId); + try { + return serializeSegment(entries.get(segmentId)); + } catch (Exception e) { + throw new java.io.IOException(e); + } + } + + @Override + public void finish(DeviceEntryDataSetHandle handle) { + finished = true; + } + } +}
diff --git a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift index d66ed10..f55f87a 100644 --- a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift +++ b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift
@@ -879,6 +879,17 @@ 2: optional string allowedUsername } +struct TFetchDeviceEntrySegmentReq { + 1: required string queryId + 2: required string planNodeId + 3: required i32 segmentId +} + +struct TFetchDeviceEntrySegmentResp { + 1: required common.TSStatus status + 2: optional binary payload +} + /** * END: Used for EXPLAIN ANALYZE **/ @@ -1440,6 +1451,10 @@ void onEndOfDataBlockEvent(TEndOfDataBlockEvent e); + TFetchDeviceEntrySegmentResp fetchDeviceEntrySegment(TFetchDeviceEntrySegmentReq req); + + common.TSStatus finishDeviceEntrySegment(1: string queryId, 2: string planNodeId); + /** Empty rpc, only for connection test */ common.TSStatus testConnectionEmptyRPC() }