This is an automated email from the ASF dual-hosted git repository.

rong pushed a commit to branch iotdb-1971
in repository https://gitbox.apache.org/repos/asf/iotdb.git

commit dd8f44d6bfca33e0e36758616f514b9ade57d022
Author: Steve Yurong Su <[email protected]>
AuthorDate: Fri Nov 26 22:32:45 2021 +0800

    UDTFJoinDataSet
---
 .../iotdb/db/query/dataset/udf/UDTFDataSet.java    |   4 +
 .../db/query/dataset/udf/UDTFJoinDataSet.java      | 111 +++++++++++++++++++++
 .../apache/iotdb/tsfile/read/common/RowRecord.java |   5 +
 3 files changed, 120 insertions(+)

diff --git 
a/server/src/main/java/org/apache/iotdb/db/query/dataset/udf/UDTFDataSet.java 
b/server/src/main/java/org/apache/iotdb/db/query/dataset/udf/UDTFDataSet.java
index 0bb6f5a..df6e5e4 100644
--- 
a/server/src/main/java/org/apache/iotdb/db/query/dataset/udf/UDTFDataSet.java
+++ 
b/server/src/main/java/org/apache/iotdb/db/query/dataset/udf/UDTFDataSet.java
@@ -126,4 +126,8 @@ public abstract class UDTFDataSet extends QueryDataSet {
   public void finalizeUDFs(long queryId) {
     udtfPlan.finalizeUDFExecutors(queryId);
   }
+
+  public UDTFPlan getUdtfPlan() {
+    return udtfPlan;
+  }
 }
diff --git 
a/server/src/main/java/org/apache/iotdb/db/query/dataset/udf/UDTFJoinDataSet.java
 
b/server/src/main/java/org/apache/iotdb/db/query/dataset/udf/UDTFJoinDataSet.java
new file mode 100644
index 0000000..9c33b04
--- /dev/null
+++ 
b/server/src/main/java/org/apache/iotdb/db/query/dataset/udf/UDTFJoinDataSet.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.dataset.udf;
+
+import org.apache.iotdb.db.exception.query.QueryProcessException;
+import org.apache.iotdb.db.utils.datastructure.TimeSelector;
+import org.apache.iotdb.tsfile.read.common.RowRecord;
+import org.apache.iotdb.tsfile.read.query.dataset.QueryDataSet;
+
+import java.io.IOException;
+
+// TODO: implements DirectAlignByTimeDataSet
+public class UDTFJoinDataSet extends QueryDataSet {
+
+  private final UDTFDataSet[] fragmentDataSets;
+
+  /**
+   * Each output column of the UDTFJoinDataSet corresponds to a two-tuple 
({@code int[]}) instance
+   * in queryDataSetOutputIndexToFragmentDataSetOutputIndex, through the 
two-tuple instance, the
+   * dataset can get the corresponding output column in {@link 
UDTFJoinDataSet#fragmentDataSets}.
+   *
+   * <p>About the two-tuple:
+   *
+   * <p>The first element is the index of the fragmentDataSet which outputs 
the column.
+   *
+   * <p>The second element is the index of the actual output column of the 
given fragmentDataSet.
+   */
+  private final int[][] resultColumnOutputIndexToFragmentDataSetOutputIndex;
+
+  private final int resultColumnsLength;
+  private final RowRecord[] rowRecordsCache;
+  private TimeSelector timeHeap;
+
+  public UDTFJoinDataSet(
+      UDTFDataSet[] fragmentDataSets, int[][] 
resultColumnOutputIndexToFragmentDataSetOutputIndex)
+      throws QueryProcessException, IOException {
+    this.fragmentDataSets = fragmentDataSets;
+    this.resultColumnOutputIndexToFragmentDataSetOutputIndex =
+        resultColumnOutputIndexToFragmentDataSetOutputIndex;
+    resultColumnsLength = 
resultColumnOutputIndexToFragmentDataSetOutputIndex.length;
+    rowRecordsCache = new RowRecord[resultColumnsLength];
+    initTimeHeap();
+  }
+
+  private void initTimeHeap() throws IOException, QueryProcessException {
+    timeHeap = new TimeSelector(resultColumnsLength << 1, true);
+    for (int i = 0; i < resultColumnsLength; ++i) {
+      UDTFDataSet fragmentDataSet = fragmentDataSets[i];
+      if (fragmentDataSet.hasNextWithoutConstraint()) {
+        rowRecordsCache[i] = fragmentDataSet.nextWithoutConstraint();
+        timeHeap.add(rowRecordsCache[i].getTimestamp());
+      }
+    }
+  }
+
+  @Override
+  public boolean hasNextWithoutConstraint() throws IOException {
+    return !timeHeap.isEmpty();
+  }
+
+  @Override
+  public RowRecord nextWithoutConstraint() throws IOException {
+    long minTime = timeHeap.pollFirst();
+    RowRecord rowRecord = new RowRecord(minTime, resultColumnsLength);
+
+    for (int i = 0; i < resultColumnsLength; ++i) {
+      int[] indexes = resultColumnOutputIndexToFragmentDataSetOutputIndex[i];
+      int fragmentDataSetIndex = indexes[0];
+      int outputColumnIndexInFragmentDataSet = indexes[1];
+
+      if (rowRecordsCache[fragmentDataSetIndex] == null) {
+        rowRecord.addField(null);
+        continue;
+      }
+
+      RowRecord fragmentRowRecord = rowRecordsCache[fragmentDataSetIndex];
+      if (fragmentRowRecord.getTimestamp() != minTime) {
+        rowRecord.addField(null);
+        continue;
+      }
+
+      
rowRecord.addField(fragmentRowRecord.getFields().get(outputColumnIndexInFragmentDataSet));
+      rowRecordsCache[fragmentDataSetIndex] = null;
+
+      if (fragmentDataSets[fragmentDataSetIndex].hasNextWithoutConstraint()) {
+        fragmentRowRecord = 
fragmentDataSets[fragmentDataSetIndex].nextWithoutConstraint();
+        rowRecordsCache[fragmentDataSetIndex] = fragmentRowRecord;
+        timeHeap.add(fragmentRowRecord.getTimestamp());
+      }
+    }
+
+    return rowRecord;
+  }
+}
diff --git 
a/tsfile/src/main/java/org/apache/iotdb/tsfile/read/common/RowRecord.java 
b/tsfile/src/main/java/org/apache/iotdb/tsfile/read/common/RowRecord.java
index 748e8ab..c764c30 100644
--- a/tsfile/src/main/java/org/apache/iotdb/tsfile/read/common/RowRecord.java
+++ b/tsfile/src/main/java/org/apache/iotdb/tsfile/read/common/RowRecord.java
@@ -38,6 +38,11 @@ public class RowRecord {
     this.fields = new ArrayList<>();
   }
 
+  public RowRecord(long timestamp, int initCapacity) {
+    this.timestamp = timestamp;
+    this.fields = new ArrayList<>(initCapacity);
+  }
+
   public RowRecord(long timestamp, List<Field> fields) {
     this.timestamp = timestamp;
     this.fields = fields;

Reply via email to