ENH: Added tasks loom/base/size and loom/base/length

src/client/plan.py

@@ -29,6 +29,8 @@ class Plan(object):
 
     TASK_BASE_GET = "loom/base/get"
     TASK_BASE_SLICE = "loom/base/slice"
+    TASK_BASE_SIZE = "loom/base/size"
+    TASK_BASE_LENGTH = "loom/base/length"
 
     TASK_DATA_CONST = "loom/data/const"
     TASK_DATA_MERGE = "loom/data/merge"
@@ -126,6 +128,20 @@ class Plan(object):
         task.mode = MODE_SIMPLE
         return self.add(task)
 
+    def task_size(self, input):
+        task = Task()
+        task.task_type = self.TASK_BASE_SIZE
+        task.inputs = (input,)
+        task.mode = MODE_SIMPLE
+        return self.add(task)
+
+    def task_length(self, input):
+        task = Task()
+        task.task_type = self.TASK_BASE_LENGTH
+        task.inputs = (input,)
+        task.mode = MODE_SIMPLE
+        return self.add(task)
+
     def task_get(self, input, index):
         task = Task()
         task.task_type = self.TASK_BASE_GET

src/libloom/tasks/basetasks.cpp

 
 #include "basetasks.h"
+#include "../data/rawdata.h"
 
 //#include "libloom/log.h"
 
@@ -24,3 +25,27 @@ void SliceTask::start(DataVector &inputs)
     auto result = input->get_slice(index[0], index[1]);
     finish(result);
 }
+
+void SizeTask::start(DataVector &inputs)
+{
+    size_t size = 0;
+    for (auto &d : inputs) {
+        size += (*d)->get_size();
+    }
+    std::shared_ptr<Data> output = std::make_shared<RawData>();
+    RawData &data = static_cast<RawData&>(*output);
+    memcpy(data.init_empty_file(worker, sizeof(size_t)), &size, sizeof(size_t));
+    finish(output);
+}
+
+void LengthTask::start(DataVector &inputs)
+{
+    size_t length = 0;
+    for (auto &d : inputs) {
+        length += (*d)->get_length();
+    }
+    std::shared_ptr<Data> output = std::make_shared<RawData>();
+    RawData &data = static_cast<RawData&>(*output);
+    memcpy(data.init_empty_file(worker, sizeof(size_t)), &length, sizeof(size_t));
+    finish(output);
+}

src/libloom/tasks/basetasks.h

@@ -18,4 +18,20 @@ public:
     void start(loom::DataVector &inputs);
 };
 
+
+class SizeTask : public loom::TaskInstance
+{
+public:
+    using TaskInstance::TaskInstance;
+    void start(loom::DataVector &inputs);
+};
+
+class LengthTask : public loom::TaskInstance
+{
+public:
+    using TaskInstance::TaskInstance;
+    void start(loom::DataVector &inputs);
+};
+
+
 #endif // LIBLOOM_TASKS_BASICTASKS_H

src/libloom/worker.cpp

@@ -84,6 +84,8 @@ void Worker::register_basic_tasks()
     // Base
     add_task_factory<GetTask>("loom/base/get");
     add_task_factory<SliceTask>("loom/base/slice");
+    add_task_factory<SizeTask>("loom/base/size");
+    add_task_factory<LengthTask>("loom/base/length");
 
     // RawData
     add_task_factory<ConstTask>("loom/data/const");

tests/client/data_test.py

 from loomenv import loom_env, LOOM_TESTPROG, LOOM_TEST_DATA_DIR  # noqa
 
+import struct
 from datetime import datetime
 import os
 
@@ -201,3 +202,20 @@ def test_split(loom_env):
     assert r2 == "Line4"
     assert r3 == "Line1\nLine2\n"
     assert r4 == ""
+
+
+def test_size_and_length(loom_env):
+    loom_env.start(1)
+
+    p = loom_env.plan()
+    text = "12345" * 5
+
+    a1 = p.task_const(text)
+    a2 = p.task_array_make((a1, a1))
+    b1 = p.task_size(a1)
+    c1 = p.task_length(a1)
+    b2 = p.task_size(a2)
+    c2 = p.task_length(a2)
+
+    u64 = struct.Struct("<Q")
+    [25, 0, 50, 2] == map(lambda x: u64.unpack(x)[0], loom_env.submit(p, (b1, c1, b2, c2)))