• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /**
2  * Copyright 2021 Huawei Technologies Co., Ltd
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #include <memory>
18 #include <string>
19 #include "common/common.h"
20 #include "minddata/dataset/engine/ir/datasetops/dataset_node.h"
21 #include "minddata/dataset/engine/ir/datasetops/skip_node.h"
22 #include "minddata/dataset/engine/ir/datasetops/take_node.h"
23 #include "minddata/dataset/engine/ir/datasetops/repeat_node.h"
24 #include "minddata/dataset/include/dataset/datasets.h"
25 
26 using namespace mindspore::dataset;
27 
28 class MindDataTestTreeModifying : public UT::DatasetOpTesting {
29  public:
30   MindDataTestTreeModifying() = default;
31 };
32 
TEST_F(MindDataTestTreeModifying,AppendChild)33 TEST_F(MindDataTestTreeModifying, AppendChild) {
34   MS_LOG(INFO) << "Doing MindDataTestTreeModifying-AppendChild";
35   /*
36   * Input tree:
37   *      ds4
38   *     /   \
39   *   ds3   ds2
40   *     |
41   *    ds1
42   *
43   * ds4->AppendChild(ds6) yields this tree
44   *
45   *      _ ds4 _
46   *     /   |   \
47   *   ds3  ds2  ds6
48   *    |
49   *   ds1
50   */
51   std::string folder_path = datasets_root_path_ + "/testPK/data/";
52   std::shared_ptr<Dataset> ds1 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
53   std::shared_ptr<Dataset> ds2 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
54   std::shared_ptr<Dataset> ds6 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
55   std::shared_ptr<Dataset> ds3 = ds1->Take(10);
56   std::shared_ptr<Dataset> ds4 = ds3->Concat({ds2});
57   Status rc;
58 
59   std::shared_ptr<DatasetNode> root = ds4->IRNode();
60   auto ir_tree = std::make_shared<TreeAdapter>();
61   rc = ir_tree->Compile(root);  // Compile adds a new RootNode to the top of the tree
62   EXPECT_EQ(rc, Status::OK());
63   // Descend two levels as Compile adds the root node and the epochctrl node on top of ds4
64   std::shared_ptr<DatasetNode> ds4_node = ir_tree->RootIRNode()->Children()[0]->Children()[0];
65   // You can inspect the plan by sending *ir_tree->RootIRNode() to std::cout
66   std::shared_ptr<DatasetNode> node_to_insert = ds6->IRNode();
67   rc = ds4_node->AppendChild(node_to_insert);
68   EXPECT_EQ(rc, Status::OK());
69   EXPECT_TRUE( ds4_node->Children()[2] == node_to_insert);
70 }
71 
TEST_F(MindDataTestTreeModifying,InsertChildAt01)72 TEST_F(MindDataTestTreeModifying, InsertChildAt01) {
73   MS_LOG(INFO) << "Doing MindDataTestTreeModifying-InsertChildAt01";
74   /*
75    * Input tree:
76    *      ds4
77    *     /   \
78    *   ds3   ds2
79    *    |     |
80    *   ds1   ds5
81    *
82    * Case 1: ds4->InsertChildAt(1, ds6) yields this tree
83    *
84    *      _ ds4 _
85    *     /   |   \
86    *   ds3  ds6  ds2
87    *    |         |
88    *   ds1       ds5
89    *
90    * Case 2: ds4->InsertChildAt(0, ds6) yields this tree
91    *
92    *      _ ds4 _
93    *     /   |   \
94    *   ds6  ds3  ds2
95    *         |    |
96    *        ds1  ds5
97    *
98    * Case 3: ds4->InsertChildAt(2, ds6) yields this tree
99    *
100    *      _ ds4 _
101    *     /   |   \
102    *   ds3  ds2  ds6
103    *    |    |
104    *   ds1  ds5
105    *
106    */
107   std::string folder_path = datasets_root_path_ + "/testPK/data/";
108   std::shared_ptr<Dataset> ds1 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
109   std::shared_ptr<Dataset> ds3 = ds1->Take(10);
110   std::shared_ptr<Dataset> ds5 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
111   std::shared_ptr<Dataset> ds2 = ds5->Repeat(4);
112   std::shared_ptr<Dataset> ds4 = ds3->Concat({ds2});
113   Status rc;
114   std::shared_ptr<DatasetNode> root = ds4->IRNode();
115   auto ir_tree = std::make_shared<TreeAdapter>();
116 
117   // Case 1:
118   rc = ir_tree->Compile(root);  // Compile adds a new RootNode to the top of the tree
119   EXPECT_EQ(rc, Status::OK());
120   // Descend two levels as Compile adds the root node and the epochctrl node on top of ds4
121   std::shared_ptr<DatasetNode> ds4_node = ir_tree->RootIRNode()->Children()[0]->Children()[0];
122   std::shared_ptr<Dataset> ds6 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
123   std::shared_ptr<DatasetNode> ds6_to_insert = ds6->IRNode();
124   std::shared_ptr<DatasetNode> ds2_node = ds4_node->Children()[1];
125   rc = ds4_node->InsertChildAt(1, ds6_to_insert);
126   EXPECT_EQ(rc, Status::OK());
127   EXPECT_TRUE( ds4_node->Children()[1] == ds6_to_insert);
128   EXPECT_TRUE( ds4_node->Children()[2] == ds2_node);
129 
130   // Case 2:
131   rc = ir_tree->Compile(root);  // Compile adds a new RootNode to the top of the tree
132   EXPECT_EQ(rc, Status::OK());
133   // Descend two levels as Compile adds the root node and the epochctrl node on top of ds4
134   ds4_node = ir_tree->RootIRNode()->Children()[0]->Children()[0];
135   ds6 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
136   ds6_to_insert = ds6->IRNode();
137   std::shared_ptr<DatasetNode> ds3_node = ds4_node->Children()[0];
138   rc = ds4_node->InsertChildAt(0, ds6_to_insert);
139   EXPECT_EQ(rc, Status::OK());
140   EXPECT_TRUE( ds4_node->Children()[0] == ds6_to_insert);
141   EXPECT_TRUE( ds4_node->Children()[1] == ds3_node);
142 
143   // Case 3:
144   rc = ir_tree->Compile(root);  // Compile adds a new RootNode to the top of the tree
145   EXPECT_EQ(rc, Status::OK());
146   // Descend two levels as Compile adds the root node and the epochctrl node on top of ds4
147   ds4_node = ir_tree->RootIRNode()->Children()[0]->Children()[0];
148   ds6 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
149   ds6_to_insert = ds6->IRNode();
150   rc = ds4_node->InsertChildAt(2, ds6_to_insert);
151   EXPECT_EQ(rc, Status::OK());
152   EXPECT_TRUE( ds4_node->Children()[2] == ds6_to_insert);
153 }
154 
TEST_F(MindDataTestTreeModifying,InsertChildAt04)155 TEST_F(MindDataTestTreeModifying, InsertChildAt04) {
156   MS_LOG(INFO) << "Doing MindDataTestTreeModifying-InsertChildAt04";
157 
158   /*
159    * Input tree:
160    *      ds4
161    *     /   \
162    *   ds3   ds2
163    *    |     |
164    *   ds1   ds5
165    *
166    */
167   std::string folder_path = datasets_root_path_ + "/testPK/data/";
168   std::shared_ptr<Dataset> ds1 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
169   std::shared_ptr<Dataset> ds3 = ds1->Take(10);
170   std::shared_ptr<Dataset> ds5 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
171   std::shared_ptr<Dataset> ds2 = ds5->Repeat(4);
172   std::shared_ptr<Dataset> ds4 = ds3->Concat({ds2});
173   Status rc;
174   std::shared_ptr<DatasetNode> root = ds4->IRNode();
175   auto ir_tree = std::make_shared<TreeAdapter>();
176 
177   // Case 4: ds4->InsertChildAt(3, ds6) raises an error
178   rc = ir_tree->Compile(root);  // Compile adds a new RootNode to the top of the tree
179   EXPECT_EQ(rc, Status::OK());
180   // Descend two levels as Compile adds the root node and the epochctrl node on top of ds4
181   std::shared_ptr<DatasetNode> ds4_node = ir_tree->RootIRNode()->Children()[0]->Children()[0];
182   std::shared_ptr<Dataset> ds6 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
183   std::shared_ptr<DatasetNode> ds6_to_insert = ds6->IRNode();
184   std::shared_ptr<DatasetNode> ds3_node = ds4_node->Children()[0];
185   std::shared_ptr<DatasetNode> ds2_node = ds4_node->Children()[1];
186   rc = ds4_node->InsertChildAt(3, ds6_to_insert);
187   EXPECT_NE(rc, Status::OK());
188   EXPECT_TRUE( ds4_node->Children()[0] == ds3_node);
189   EXPECT_TRUE( ds4_node->Children()[1] == ds2_node);
190 
191   // Case 5: ds4->InsertChildAt(-1, ds6) raises an error
192   rc = ir_tree->Compile(root);  // Compile adds a new RootNode to the top of the tree
193   EXPECT_EQ(rc, Status::OK());
194   // Descend two levels as Compile adds the root node and the epochctrl node on top of ds4
195   ds4_node = ir_tree->RootIRNode()->Children()[0]->Children()[0];
196   ds6 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
197   ds6_to_insert = ds6->IRNode();
198   ds3_node = ds4_node->Children()[0];
199   ds2_node = ds4_node->Children()[1];
200   rc = ds4_node->InsertChildAt(-1, ds6_to_insert);
201   EXPECT_NE(rc, Status::OK());
202   EXPECT_TRUE( ds4_node->Children()[0] == ds3_node);
203   EXPECT_TRUE( ds4_node->Children()[1] == ds2_node);
204 }
205 
TEST_F(MindDataTestTreeModifying,InsertAbove01)206 TEST_F(MindDataTestTreeModifying, InsertAbove01) {
207   MS_LOG(INFO) << "Doing MindDataTestTreeModifying-InsertAbove01";
208   /*
209    * Insert the input <node> above this node
210    * Input tree:
211    *       ds4
212    *      /   \
213    *     ds3  ds2
214    *      |
215    *     ds1
216    *
217    * Case 1: If we want to insert a new node ds5 between ds4 and ds3, use
218    *           ds3->InsertAbove(ds5)
219    *
220    *       ds4
221    *      /   \
222    *     ds5  ds2
223    *      |
224    *     ds3
225    *      |
226    *     ds1
227    *
228    * Case 2: Likewise, ds2->InsertAbove(ds6) yields
229    *
230    *       ds4
231    *      /   \
232    *     ds3  ds6
233    *      |    |
234    *     ds1  ds2
235    *
236    * Case 3: We can insert a new node between ds3 and ds1 by ds1->InsertAbove(ds7)
237    *
238    *       ds4
239    *      /   \
240    *     ds3  ds2
241    *      |
242    *     ds7
243    *      |
244    *     ds1
245    *
246    */
247   // Case 1
248   std::string folder_path = datasets_root_path_ + "/testPK/data/";
249   std::shared_ptr<Dataset> ds1 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
250   std::shared_ptr<Dataset> ds2 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
251   std::shared_ptr<Dataset> ds3 = ds1->Take(10);
252   std::shared_ptr<Dataset> ds4 = ds3->Concat({ds2});
253   Status rc;
254 
255   std::shared_ptr<DatasetNode> root = ds4->IRNode();
256   auto ir_tree = std::make_shared<TreeAdapter>();
257   rc = ir_tree->Compile(root);  // Compile adds a new RootNode to the top of the tree
258   EXPECT_EQ(rc, Status::OK());
259   // Descend two levels as Compile adds the root node and the epochctrl node on top of ds4
260   std::shared_ptr<DatasetNode> ds4_node = ir_tree->RootIRNode()->Children()[0]->Children()[0];
261   std::shared_ptr<DatasetNode> ds3_node = ds4_node->Children()[0];
262   std::shared_ptr<SkipNode> ds5_to_insert = std::make_shared<SkipNode>(nullptr, 1);
263   rc = ds3_node->InsertAbove(ds5_to_insert);
264   EXPECT_EQ(rc, Status::OK());
265   EXPECT_TRUE(ds5_to_insert->Children()[0] == ds3_node);
266   EXPECT_TRUE( ds4_node->Children()[0] == ds5_to_insert);
267 }
268 
TEST_F(MindDataTestTreeModifying,InsertAbove02)269 TEST_F(MindDataTestTreeModifying, InsertAbove02) {
270   MS_LOG(INFO) << "Doing MindDataTestTreeModifying-InsertAbove02";
271 
272   // Case 2
273   std::string folder_path = datasets_root_path_ + "/testPK/data/";
274   std::shared_ptr<Dataset> ds1 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
275   std::shared_ptr<Dataset> ds2 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
276   std::shared_ptr<Dataset> ds3 = ds1->Take(10);
277   std::shared_ptr<Dataset> ds4 = ds3 + ds2;
278   Status rc;
279 
280   std::shared_ptr<DatasetNode> root = ds4->IRNode();
281   auto ir_tree = std::make_shared<TreeAdapter>();
282   rc = ir_tree->Compile(root);  // Compile adds a new RootNode to the top of the tree
283   EXPECT_EQ(rc, Status::OK());
284   // Descend two levels as Compile adds the root node and the epochctrl node on top of ds4
285   std::shared_ptr<DatasetNode> ds4_node = ir_tree->RootIRNode()->Children()[0]->Children()[0];
286   std::shared_ptr<DatasetNode> ds2_node = ds4_node->Children()[1];
287   std::shared_ptr<TakeNode> ds6_to_insert = std::make_shared<TakeNode>(nullptr, 12);
288   rc = ds2_node->InsertAbove(ds6_to_insert);
289   EXPECT_EQ(rc, Status::OK());
290   EXPECT_TRUE(ds6_to_insert->Children()[0] == ds2_node);
291   EXPECT_TRUE( ds4_node->Children()[1] == ds6_to_insert);
292 }
293 
TEST_F(MindDataTestTreeModifying,InsertAbove03)294 TEST_F(MindDataTestTreeModifying, InsertAbove03) {
295   MS_LOG(INFO) << "Doing MindDataTestTreeModifying-InsertAbove03";
296 
297   // Case 3
298   std::string folder_path = datasets_root_path_ + "/testPK/data/";
299   std::shared_ptr<Dataset> ds1 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
300   std::shared_ptr<Dataset> ds2 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
301   std::shared_ptr<Dataset> ds3 = ds1->Take(10);
302   std::shared_ptr<Dataset> ds4 = ds3->Concat({ds2});
303   Status rc;
304 
305   std::shared_ptr<DatasetNode> root = ds4->IRNode();
306   auto ir_tree = std::make_shared<TreeAdapter>();
307   rc = ir_tree->Compile(root);  // Compile adds a new RootNode to the top of the tree
308   EXPECT_EQ(rc, Status::OK());
309   // Descend two levels as Compile adds the root node and the epochctrl node on top of ds4
310   std::shared_ptr<DatasetNode> ds4_node = ir_tree->RootIRNode()->Children()[0]->Children()[0];
311   std::shared_ptr<DatasetNode> ds3_node = ds4_node->Children()[0];
312   std::shared_ptr<DatasetNode> ds1_node = ds3_node->Children()[0];
313   std::shared_ptr<RepeatNode> ds7_to_insert = std::make_shared<RepeatNode>(nullptr, 3);
314   rc = ds1_node->InsertAbove(ds7_to_insert);
315   EXPECT_TRUE(ds7_to_insert->Children()[0] == ds1_node);
316   EXPECT_TRUE( ds3_node->Children()[0] == ds7_to_insert);
317 }
318 
TEST_F(MindDataTestTreeModifying,Drop01)319 TEST_F(MindDataTestTreeModifying, Drop01) {
320   MS_LOG(INFO) << "Doing MindDataTestTreeModifying-Drop01";
321   /*
322    * Drop() detaches this node from the tree it is in. Calling Drop() from a standalone node is a no-op.
323    *
324    * Input tree:
325    *       ds10
326    *      /    \
327    *    ds9    ds6
328    *     |   /  |  \
329    *    ds8 ds5 ds4 ds1
330    *     |     /  \
331    *    ds7  ds3  ds2
332    *
333    * Case 1: When the node has no child and no sibling, Drop() detaches the node from its tree.
334    *
335    *   ds7->Drop() yields the tree below:
336    *
337    *       ds10
338    *      /    \
339    *    ds9    ds6
340    *     |   /  |  \
341    *    ds8 ds5 ds4 ds1
342    *           /  \
343    *         ds3  ds2
344    *
345    * Case 2: When the node has one child and no sibling, Drop() detaches the node from its tree and the node's child
346    *         becomes its parent's child.
347    *
348    *   ds8->Drop() yields the tree below:
349    *
350    *       ds10
351    *      /    \
352    *    ds9    ds6
353    *     |   /  |  \
354    *    ds7 ds5 ds4 ds1
355    *           /  \
356    *         ds3  ds2
357    *
358    */
359   std::string folder_path = datasets_root_path_ + "/testPK/data/";
360   std::shared_ptr<Dataset> ds7 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
361   std::shared_ptr<Dataset> ds8 = ds7->Take(20);
362   std::shared_ptr<Dataset> ds9 = ds8->Skip(1);
363   std::shared_ptr<Dataset> ds3 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
364   std::shared_ptr<Dataset> ds2 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
365   std::shared_ptr<Dataset> ds4 = ds3->Concat({ds2});
366   std::shared_ptr<Dataset> ds6 = ds4->Take(13);
367   std::shared_ptr<Dataset> ds10 = ds9 + ds6;
368   Status rc;
369 
370   std::shared_ptr<DatasetNode> root = ds10->IRNode();
371   auto ir_tree = std::make_shared<TreeAdapter>();
372 
373   // Case 1
374   rc = ir_tree->Compile(root);  // Compile adds a new RootNode to the top of the tree
375   EXPECT_EQ(rc, Status::OK());
376   // Descend two levels as Compile adds the root node and the epochctrl node on top of ds4
377   std::shared_ptr<DatasetNode> ds10_node = ir_tree->RootIRNode()->Children()[0]->Children()[0];
378   std::shared_ptr<DatasetNode> ds9_node = ds10_node->Children()[0];
379   std::shared_ptr<DatasetNode> ds8_node = ds9_node->Children()[0];
380   std::shared_ptr<DatasetNode> ds7_node = ds8_node->Children()[0];
381   rc = ds7_node->Drop();
382   EXPECT_EQ(rc, Status::OK());
383   // ds8 becomes a childless node
384   EXPECT_TRUE(ds8_node->Children().empty());
385   EXPECT_TRUE(ds7_node->Children().empty());
386 
387   // Case 2
388   rc = ir_tree->Compile(root);  // Compile adds a new RootNode to the top of the tree
389   EXPECT_EQ(rc, Status::OK());
390   // Descend two levels as Compile adds the root node and the epochctrl node on top of ds4
391   ds10_node = ir_tree->RootIRNode()->Children()[0]->Children()[0];
392   ds9_node = ds10_node->Children()[0];
393   ds8_node = ds9_node->Children()[0];
394   ds7_node = ds8_node->Children()[0];
395   rc = ds8_node->Drop();
396   EXPECT_EQ(rc, Status::OK());
397   // ds7 becomes a child of ds9
398   EXPECT_TRUE(ds9_node->Children()[0] == ds7_node);
399   EXPECT_TRUE(ds8_node->Children().empty());
400 }
401 
TEST_F(MindDataTestTreeModifying,Drop03)402 TEST_F(MindDataTestTreeModifying, Drop03) {
403   MS_LOG(INFO) << "Doing MindDataTestTreeModifying-Drop03";
404   /* Case 3: When the node has more than one child and no sibling, Drop() detaches the node from its tree and the node's
405    *         children become its parent's children.
406    *
407    *   When the input tree is
408    *       ds10
409    *      /    \
410    *    ds9    ds6
411    *     |      |
412    *    ds8    ds4
413    *     |    /   \
414    *    ds7  ds3  ds2
415    *
416    *
417    *   ds4->Drop() will raise an error because we cannot add the children of an n-ary operator (ds4) to a unary operator
418    *   (ds6).
419    *
420    */
421   std::string folder_path = datasets_root_path_ + "/testPK/data/";
422   std::shared_ptr<Dataset> ds7 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
423   std::shared_ptr<Dataset> ds8 = ds7->Take(20);
424   std::shared_ptr<Dataset> ds9 = ds8->Skip(1);
425   std::shared_ptr<Dataset> ds3 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
426   std::shared_ptr<Dataset> ds2 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
427   std::shared_ptr<Dataset> ds4 = ds3->Concat({ds2});
428   std::shared_ptr<Dataset> ds6 = ds4->Take(13);
429   std::shared_ptr<Dataset> ds10 = ds9 + ds6;
430   Status rc;
431 
432   std::shared_ptr<DatasetNode> root = ds10->IRNode();
433   auto ir_tree = std::make_shared<TreeAdapter>();
434   rc = ir_tree->Compile(root);  // Compile adds a new RootNode to the top of the tree
435   EXPECT_EQ(rc, Status::OK());
436   // Descend two levels as Compile adds the root node and the epochctrl node on top of ds4
437   std::shared_ptr<DatasetNode> ds10_node = ir_tree->RootIRNode()->Children()[0]->Children()[0];
438   std::shared_ptr<DatasetNode> ds6_node = ds10_node->Children()[1];
439   std::shared_ptr<DatasetNode> ds4_node = ds6_node->Children()[0];
440   std::shared_ptr<DatasetNode> ds3_node = ds4_node->Children()[0];
441   std::shared_ptr<DatasetNode> ds2_node = ds4_node->Children()[1];
442   rc = ds4_node->Drop();
443   EXPECT_NE(rc, Status::OK());
444 }
445 
TEST_F(MindDataTestTreeModifying,Drop04)446 TEST_F(MindDataTestTreeModifying, Drop04) {
447   MS_LOG(INFO) << "Doing MindDataTestTreeModifying-Drop04";
448   /* Case 4: When the node has no child but has siblings, Drop() detaches the node from its tree and its siblings will be
449    *         squeezed left.
450    *
451    * Input tree:
452    *       ds10
453    *      /    \
454    *    ds9    ds6
455    *     |   /  |  \
456    *    ds8 ds5 ds4 ds1
457    *     |     /  \
458    *    ds7  ds3  ds2
459    *
460    *   ds5->Drop() yields the tree below:
461    *
462    *       ds10
463    *      /    \
464    *    ds9    ds6
465    *     |     /  \
466    *    ds8   ds4 ds1
467    *     |    /  \
468    *    ds7 ds3  ds2
469    *
470    */
471   std::string folder_path = datasets_root_path_ + "/testPK/data/";
472   std::shared_ptr<Dataset> ds7 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
473   std::shared_ptr<Dataset> ds8 = ds7->Take(20);
474   std::shared_ptr<Dataset> ds9 = ds8->Skip(1);
475   std::shared_ptr<Dataset> ds3 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
476   std::shared_ptr<Dataset> ds2 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
477   std::shared_ptr<Dataset> ds4 = ds3->Concat({ds2});
478   std::shared_ptr<Dataset> ds5 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
479   std::shared_ptr<Dataset> ds1 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
480   std::shared_ptr<Dataset> ds6 = ds5->Concat({ds4, ds1});
481   std::shared_ptr<Dataset> ds10 = ds9 + ds6;
482   Status rc;
483 
484   std::shared_ptr<DatasetNode> root = ds10->IRNode();
485   auto ir_tree = std::make_shared<TreeAdapter>();
486   rc = ir_tree->Compile(root);  // Compile adds a new RootNode to the top of the tree
487   EXPECT_EQ(rc, Status::OK());
488   // Descend two levels as Compile adds the root node and the epochctrl node on top of ds4
489   std::shared_ptr<DatasetNode> ds10_node = ir_tree->RootIRNode()->Children()[0]->Children()[0];
490   std::shared_ptr<DatasetNode> ds6_node = ds10_node->Children()[1];
491   std::shared_ptr<DatasetNode> ds5_node = ds6_node->Children()[0];
492   std::shared_ptr<DatasetNode> ds4_node = ds6_node->Children()[1];
493   EXPECT_TRUE(ds5_node->IsDataSource());
494   EXPECT_TRUE(ds6_node->IsNaryOperator());
495   rc = ds5_node->Drop();
496   EXPECT_EQ(rc, Status::OK());
497   EXPECT_TRUE(ds6_node->Children().size() == 2);
498   EXPECT_TRUE(ds6_node->Children()[0] == ds4_node);
499   EXPECT_TRUE(ds5_node->Children().empty());
500 }
501 
TEST_F(MindDataTestTreeModifying,Drop05)502 TEST_F(MindDataTestTreeModifying, Drop05) {
503   MS_LOG(INFO) << "Doing MindDataTestTreeModifying-Drop05";
504   /*
505    * Case 5: When the node has only one child but has siblings, Drop() detaches the node from its tree and the node's
506    *         children become its parent's children.
507    *
508    * Input tree:
509    *       ds10
510    *      /    \
511    *    ds9    ds6
512    *     |   /  |  \
513    *    ds8 ds5 ds4 ds1
514    *     |      |
515    *    ds7    ds3
516    *
517    *   ds4->Drop() yields the tree below:
518    *
519    *       ds10
520    *      /    \
521    *    ds9    ds6
522    *     |   /  |  \
523    *    ds8 ds5 ds3 ds1
524    *     |
525    *    ds7
526    *
527    */
528   std::string folder_path = datasets_root_path_ + "/testPK/data/";
529   std::shared_ptr<Dataset> ds7 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
530   std::shared_ptr<Dataset> ds8 = ds7->Take(20);
531   std::shared_ptr<Dataset> ds9 = ds8->Skip(1);
532   std::shared_ptr<Dataset> ds3 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
533   std::shared_ptr<Dataset> ds4 = ds3->Skip(1);
534   std::shared_ptr<Dataset> ds5 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
535   std::shared_ptr<Dataset> ds1 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
536   std::shared_ptr<Dataset> ds6 = ds5->Concat({ds4, ds1});
537   std::shared_ptr<Dataset> ds10 = ds9 + ds6;
538   Status rc;
539 
540   std::shared_ptr<DatasetNode> root = ds10->IRNode();
541   auto ir_tree = std::make_shared<TreeAdapter>();
542   rc = ir_tree->Compile(root);  // Compile adds a new RootNode to the top of the tree
543   EXPECT_EQ(rc, Status::OK());
544   // Descend two levels as Compile adds the root node and the epochctrl node on top of ds4
545   std::shared_ptr<DatasetNode> ds10_node = ir_tree->RootIRNode()->Children()[0]->Children()[0];
546   std::shared_ptr<DatasetNode> ds6_node = ds10_node->Children()[1];
547   std::shared_ptr<DatasetNode> ds4_node = ds6_node->Children()[1];
548   std::shared_ptr<DatasetNode> ds3_node = ds4_node->Children()[0];
549   rc = ds4_node->Drop();
550   EXPECT_EQ(rc, Status::OK());
551   EXPECT_TRUE(ds6_node->Children().size() == 3);
552   EXPECT_TRUE(ds6_node->Children()[1] == ds3_node);
553   EXPECT_TRUE(ds4_node->Children().empty());
554 }
555 
TEST_F(MindDataTestTreeModifying,Drop06)556 TEST_F(MindDataTestTreeModifying, Drop06) {
557   MS_LOG(INFO) << "Doing MindDataTestTreeModifying-Drop06";
558   /*
559    * Case 6: When the node has more than one child and more than one sibling, Drop() will raise an error.
560    *         If we want to drop ds4 from the input tree, ds4->Drop() will not work. We will have to do it
561    *         with a combination of Drop(), InsertChildAt()
562    *
563    * Input tree:
564    *       ds10
565    *      /    \
566    *    ds9    ds6
567    *     |   /  |  \
568    *    ds8 ds5 ds4 ds1
569    *     |     /  \
570    *    ds7  ds3  ds2
571    *
572    * If we want to form this tree below:
573    *
574    *       ds10
575    *      /    \
576    *    ds9    ds6_____
577    *     |   /  |   |  \
578    *    ds8 ds5 ds3 ds2 ds1
579    *     |
580    *    ds7
581    *
582    */
583   std::string folder_path = datasets_root_path_ + "/testPK/data/";
584   std::shared_ptr<Dataset> ds7 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
585   std::shared_ptr<Dataset> ds8 = ds7->Take(20);
586   std::shared_ptr<Dataset> ds9 = ds8->Skip(1);
587   std::shared_ptr<Dataset> ds3 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
588   std::shared_ptr<Dataset> ds2 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
589   std::shared_ptr<Dataset> ds4 = ds3->Concat({ds2});
590   std::shared_ptr<Dataset> ds5 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
591   std::shared_ptr<Dataset> ds1 = ImageFolder(folder_path, false, std::make_shared<SequentialSampler>(0, 11));
592   std::shared_ptr<Dataset> ds6 = ds5->Concat({ds4, ds1});  // ds1 is put after (ds5, ds4)!!!
593   std::shared_ptr<Dataset> ds10 = ds9 + ds6;
594   Status rc;
595 
596   std::shared_ptr<DatasetNode> root = ds10->IRNode();
597   auto ir_tree = std::make_shared<TreeAdapter>();
598   rc = ir_tree->Compile(root);  // Compile adds a new RootNode to the top of the tree
599   EXPECT_EQ(rc, Status::OK());
600   // Descend two levels as Compile adds the root node and the epochctrl node on top of ds4
601   std::shared_ptr<DatasetNode> ds10_node = ir_tree->RootIRNode()->Children()[0]->Children()[0];
602   std::shared_ptr<DatasetNode> ds6_node = ds10_node->Children()[1];
603   std::shared_ptr<DatasetNode> ds4_node = ds6_node->Children()[1];
604   rc = ds4_node->Drop();
605   EXPECT_NE(rc, Status::OK());
606 }
607