tuplex · KorlaMarch · Jan 22, 2022 · Jan 28, 2022 · Jan 28, 2022 · Feb 11, 2022
diff --git a/tuplex/core/include/DataSet.h b/tuplex/core/include/DataSet.h
@@ -125,10 +125,10 @@ namespace tuplex {
 
         /*!
          * action that displays tuples as nicely formatted table
-         * @param numRows how many rows to print, i.e. top numRows are printed.xs
+         * @param numRows how many rows to print, i.e. top numRows are printed.xs, -1 means print all rows
          * @param os ostream where to print table to
          */
-        virtual void show(const int64_t numRows = -1, std::ostream &os = std::cout);
+        virtual void show(int64_t numRows = -1, std::ostream &os = std::cout);
 
         // named dataset management functions
         /*!
@@ -252,22 +252,49 @@ namespace tuplex {
          * @param memoryLayout
          * @return
          */
-        virtual DataSet& cache(const Schema::MemoryLayout& memoryLayout, bool storeSpecialized);
-        DataSet& cache(bool storeSpecialized=true) { return cache(Schema::MemoryLayout::ROW, storeSpecialized); }
+        virtual DataSet &cache(const Schema::MemoryLayout &memoryLayout, bool storeSpecialized);
+
+        DataSet &cache(bool storeSpecialized = true) { return cache(Schema::MemoryLayout::ROW, storeSpecialized); }
 
         /*!
          * helper setter without checks, to update internal column names.
          */
         void setColumns(const std::vector<std::string> &columnNames) { _columnNames = columnNames; }
 
-        // these are actions that cause execution
+        /*!
+         * Execute the pipeline and return all outputs
+         * @param os the logging output
+         * @return the output of the execution
+         */
         virtual std::shared_ptr<ResultSet> collect(std::ostream &os = std::cout);
 
-        virtual std::shared_ptr<ResultSet> take(int64_t numElements, std::ostream &os = std::cout);
+        /*!
+         * Execute the pipeline and take a subset of the output from the top and bottom rows.
+         * If both top and bottom rows limit exist, then the top and bottom rows will be concatenated.
+         * In the case where topLimit + bottomLimit exceeds the output size, all rows will be taken.
+         * To take all rows, pass in either topLimit=size_t::max(), bottomLimit=size_t::max(), or both.
+         * @param topLimit number of top rows to take. size_t::max() means taking all rows
+         * @param bottomLimit number of bottom rows to take. size_t::max() means taking all rows
+         * @param os the logging output
+         * @return result of the execution, trim to the size of top and bottom limit.
+         */
+        virtual std::shared_ptr<ResultSet> take(size_t topLimit, size_t bottomLimit, std::ostream &os = std::cout);
 
+        /*!
+         * Execute the pipeline and return all outputs as vector
+         * @param os the logging output
+         * @return the output of the execution in vector
+         */
         virtual std::vector<Row> collectAsVector(std::ostream &os = std::cout);
 
-        virtual std::vector<Row> takeAsVector(int64_t numElements, std::ostream &os = std::cout);
+        /*!
+         * Execute the pipeline and take a subset of the output from the top rows, return as vector
+         * In the case where numElements exceeds the output size, all rows will be taken.
+         * @param numElements number of top rows to take. size_t::max() means taking all rows
+         * @param os the logging output
+         * @return result of the execution in vector, trim to the size of numElements
+         */
+        virtual std::vector<Row> takeAsVector(size_t numElements, std::ostream &os = std::cout);
 
         /*!
          * saves dataset to file. There are multiple options to control the behavior

diff --git a/tuplex/core/include/EmptyDataset.h b/tuplex/core/include/EmptyDataset.h
@@ -67,18 +67,20 @@ namespace tuplex {
         virtual DataSet& aggregateByKey(const UDF& aggCombine, const UDF& aggUDF, const Row& aggInitial, const std::vector<std::string> &keyColumns) override { return *this; }
 
         //virtual void show(const int64_t numRows=-1, std::ostream& os=std::cout) override;
-        virtual std::shared_ptr<ResultSet> collect(std::ostream& os) override;
+        virtual std::shared_ptr<ResultSet> collect(std::ostream &os) override;
 
         // take / collect will print out the error only
-        virtual std::shared_ptr<ResultSet> take(int64_t numElements, std::ostream& os) override;
+        virtual std::shared_ptr<ResultSet> take(size_t topLimit, size_t bottomLimit, std::ostream &os) override;
 
         //virtual void show(const int64_t numRows=-1, std::ostream& os=std::cout) override;
-        virtual std::vector<Row> collectAsVector(std::ostream& os) override;
+        virtual std::vector<Row> collectAsVector(std::ostream &os) override;
 
-        // take / collect will print out the error only
-        virtual std::vector<Row> takeAsVector(int64_t numElements, std::ostream& os) override;
+        /*!
+         * take / collect will print out the error only, return empty rows
+         */
+        virtual std::vector<Row> takeAsVector(size_t numElements, std::ostream &os) override;
 
-        DataSet& cache(const Schema::MemoryLayout& memoryLayout, bool storeSpecialized) override {
+        DataSet &cache(const Schema::MemoryLayout &memoryLayout, bool storeSpecialized) override {
             return *this;
         }
     };

diff --git a/tuplex/core/include/ErrorDataSet.h b/tuplex/core/include/ErrorDataSet.h
@@ -90,13 +90,13 @@ namespace tuplex {
         std::shared_ptr<ResultSet> collect(std::ostream& os) override;
 
         // take / collect will print out the error only
-        std::shared_ptr<ResultSet> take(int64_t numElements, std::ostream& os) override;
+        std::shared_ptr<ResultSet> take(size_t topLimit, size_t bottomLimit, std::ostream& os) override;
 
         //virtual void show(const int64_t numRows=-1, std::ostream& os=std::cout) override;
         std::vector<Row> collectAsVector(std::ostream& os) override;
 
         // take / collect will print out the error only
-        std::vector<Row> takeAsVector(int64_t numElements, std::ostream& os) override;
+        std::vector<Row> takeAsVector(size_t numElements, std::ostream& os) override;
     };
 }
 

diff --git a/tuplex/core/include/Executor.h b/tuplex/core/include/Executor.h
@@ -48,8 +48,9 @@ namespace tuplex {
         ExecutorTaskQueueType _queue;
         std::mutex _completedTasksMutex;
         std::vector<IExecutorTask*> _completedTasks;
-        std::atomic_int _numPendingTasks;
-        std::atomic_int _numCompletedTasks;
+        std::atomic_int _numPendingTasks{};
+        std::atomic_int _numCompletedTasks{};
+
     public:
 
         WorkQueue();

diff --git a/tuplex/core/include/LocalEngine.h b/tuplex/core/include/LocalEngine.h
@@ -16,7 +16,28 @@
 #include <TSingleton.h>
 #include "RESTInterface.h"
 
+
 namespace tuplex {
+    struct ExecutorConfig {
+        size_t _size; // size in bytes that each executor should have
+        size_t _blockSize; // size of individual blocks used (can be used for coarse or fine grained parallelism)
+        size_t _runTimeMemory;
+        size_t _runTimeMemoryDefaultBlockSize;
+        URI _cache_path;
+
+        bool operator==(const ExecutorConfig &rhs) const {
+            return _size == rhs._size &&
+                   _blockSize == rhs._blockSize &&
+                   _runTimeMemory == rhs._runTimeMemory &&
+                   _runTimeMemoryDefaultBlockSize == rhs._runTimeMemoryDefaultBlockSize &&
+                   _cache_path == rhs._cache_path;
+        }
+
+        bool operator!=(const ExecutorConfig &rhs) const {
+            return !(rhs == *this);
+        }
+    };
+
     /*!
      * local execution engine. Provides local executors for a context
      * THIS IS NOT THREADSAFE. Should be only accessed by driver thread.
@@ -25,16 +46,18 @@ namespace tuplex {
 
     private:
         // non-detached executor that serves as the driver
-        std::unique_ptr<Executor> _driver;
+        std::shared_ptr<Executor> _driver;
+        ExecutorConfig _driver_cfg;
 
         std::vector<std::unique_ptr<Executor>> _executors;
-        std::map<Executor*, size_t> _refCounts; //! reference counts for each executor
+        std::map<Executor *, size_t> _refCounts; //! reference counts for each executor
+
+        LocalEngine(const LocalEngine &);
 
-        LocalEngine(const LocalEngine&);
-        void operator = (const LocalEngine&);
+        void operator=(const LocalEngine &);
 
         // The local task queue
-        WorkQueue  _queue;
+        WorkQueue _queue;
 
     protected:
         LocalEngine();
@@ -63,33 +86,33 @@ namespace tuplex {
          * @param cache_path directory where subfolders will be created for all executors to be started
          * @return array of executor references
          */
-        std::vector<Executor*> getExecutors(const size_t num,
-                                            const size_t size,
-                                            const size_t blockSize,
-                                            const size_t runTimeMemory,
-                                            const size_t runTimeMemoryDefaultBlockSize,
-                const URI& cache_path);
+        std::vector<Executor *> getExecutors(const size_t num,
+                                             const size_t size,
+                                             const size_t blockSize,
+                                             const size_t runTimeMemory,
+                                             const size_t runTimeMemoryDefaultBlockSize,
+                                             const URI &cache_path);
 
         /*!
          * releases executors (invoked by context)
          * @param executors
          * @param ctx
          */
-        void freeExecutors(const std::vector<Executor*>& executors, const Context* ctx=nullptr);
+        void freeExecutors(const std::vector<Executor *> &executors, const Context *ctx = nullptr);
 
-        Executor* getDriver(const size_t size,
-                            const size_t blockSize,
-                            const size_t runTimeMemory,
-                            const size_t runTimeMemoryDefaultBlockSize,
-                            const URI& cache_path);
+        std::shared_ptr<Executor> getDriver(const size_t size,
+                                            const size_t blockSize,
+                                            const size_t runTimeMemory,
+                                            const size_t runTimeMemoryDefaultBlockSize,
+                                            const URI &cache_path);
 
         void release();
 
         /*!
          * retrieves the global work queue for local executors
          * @return
          */
-        WorkQueue& getQueue() { return _queue; }
+        WorkQueue &getQueue() { return _queue; }
     };
 }
 #endif //TUPLEX_LOCALENGINE_H
diff --git a/tuplex/core/include/Partition.h b/tuplex/core/include/Partition.h
@@ -157,7 +157,7 @@ namespace tuplex {
          * return how much capacity is left, i.e. how many bytes can be actually written
          * @return
          */
-        size_t capacity() { return _size - sizeof(int64_t); }
+        size_t capacity() const { return _size - sizeof(int64_t); }
 
         uniqueid_t uuid() const { return _uuid; }
 

diff --git a/tuplex/core/include/PartitionUtils.h b/tuplex/core/include/PartitionUtils.h
@@ -0,0 +1,46 @@
+//--------------------------------------------------------------------------------------------------------------------//
+//                                                                                                                    //
+//                                      Tuplex: Blazing Fast Python Data Science                                      //
+//                                                                                                                    //
+//                                                                                                                    //
+//  (c) 2017 - 2021, Tuplex team                                                                                      //
+//  Created by March Boonyapaluk first on 4/19/2021                                                                   //
+//  License: Apache 2.0                                                                                               //
+//--------------------------------------------------------------------------------------------------------------------//
+
+#ifndef TUPLEX_PARTITIONUTILS_H
+#define TUPLEX_PARTITIONUTILS_H
+
+#include <vector>
+#include <physical/TransformStage.h>
+#include <Executor.h>
+
+namespace tuplex {
+    /*!
+     * Trim list of partitions so that it includes up to the first n rows and the last m rows
+     * if n + m > number of rows in input partitions, the partitions will remain unchanged
+     * @param partitions [in,out] the list of partitions to trim
+     * @param topLimit n, the number of top rows to include
+     * @param bottomLimit m, the number of bottom rows to include
+     * @param tstage pointer to transform stage, might be used to generate new partition
+     * @param exec pointer to executor, might be used to allocate new partition
+     */
+    void trimPartitionsToLimit(std::vector<Partition *> &partitions, size_t topLimit, size_t bottomLimit,
+                               TransformStage *tstage, Executor *exec);
+
+    /*!
+     * Create a newly allocated partition with the same data as the specified partition, but with the first n rows removed
+     * @param p_in the input partition
+     * @param numToSkip number of rows to remove from the new partition
+     * @param tstage pointer to transform stage, used to generate new partition
+     * @param exec pointer to executor, used to allocate new partition
+     * @return the new partition
+     */
+    Partition *newPartitionWithSkipRows(Partition *p_in,
+                                        size_t numToSkip,
+                                        TransformStage *tstage,
+                                        Executor *exec);
+
+}
+
+#endif //TUPLEX_PARTITIONUTILS_H
diff --git a/tuplex/core/include/ee/IBackend.h b/tuplex/core/include/ee/IBackend.h
@@ -29,19 +29,22 @@ namespace tuplex {
     class IBackend {
     public:
         IBackend() = delete;
-        IBackend(const IBackend& other) = delete;
-        IBackend(const Context& context) : _context(context) {}
+
+        IBackend(const IBackend &other) = delete;
+
+        IBackend(const Context &context) : _context(context) {}
 
         // driver, i.e. where to store local data.
-        virtual Executor* driver() = 0;
-        virtual void execute(PhysicalStage* stage) = 0;
+        virtual Executor *driver() = 0;
+
+        virtual void execute(PhysicalStage *stage) = 0;
 
         virtual ~IBackend() {} // virtual destructor needed b.c. of smart pointers
 
-        virtual const Context& context() const { return _context; }
+        virtual const Context &context() const { return _context; }
 
     private:
-        const Context& _context;
+        const Context &_context;
     };
 
     inline std::unordered_map<std::tuple<int64_t, ExceptionCode>, size_t> merge_ecounts(std::unordered_map<std::tuple<int64_t, ExceptionCode>, size_t> lhs,

diff --git a/tuplex/core/include/ee/local/LocalBackend.h b/tuplex/core/include/ee/local/LocalBackend.h
@@ -40,14 +40,15 @@ namespace tuplex {
          * constructor for convenience
          * @param context
          */
-        explicit LocalBackend(const Context& context);
+        explicit LocalBackend(const Context &context);
 
-        Executor* driver() override; // for local execution
+        Executor *driver() override; // for local execution
+
+        void execute(PhysicalStage *stage) override;
 
-        void execute(PhysicalStage* stage) override;
     private:
-        Executor *_driver; //! driver from local backend...
-        std::vector<Executor*> _executors; //! drivers to be used
+        std::shared_ptr<Executor> _driver; //! driver from local backend...
+        std::vector<Executor *> _executors; //! drivers to be used
         std::unique_ptr<JITCompiler> _compiler;
 
         HistoryServerConnection _historyConn;

diff --git a/tuplex/core/include/logical/TakeOperator.h b/tuplex/core/include/logical/TakeOperator.h
@@ -17,15 +17,16 @@
 namespace tuplex {
     class TakeOperator : public LogicalOperator {
     private:
-        int64_t _limit;
+        size_t _topLimit;
+        size_t _bottomLimit;
     public:
         LogicalOperator *clone() override;
 
     public:
-        TakeOperator(LogicalOperator *parent, const int64_t numElements);
+        TakeOperator(LogicalOperator *parent, size_t topLimit, size_t bottomLimit);
 
         std::string name() override {
-            if(_limit < 0 || std::numeric_limits<int64_t>::max() == _limit)
+            if(_topLimit == std::numeric_limits<size_t>::max() || _bottomLimit == std::numeric_limits<size_t>::max())
                 return "collect";
             return "take";
         }
@@ -37,8 +38,9 @@ namespace tuplex {
 
         bool good() const override;
 
-        int64_t limit() { return _limit; }
+        size_t topLimit() const { return _topLimit; }
 
+        size_t bottomLimit() const { return _bottomLimit; }
 
         std::vector<Row> getSample(const size_t num) const override;
 

diff --git a/tuplex/core/include/physical/StageBuilder.h b/tuplex/core/include/physical/StageBuilder.h
@@ -76,8 +76,9 @@ namespace tuplex {
             void addFileInput(FileInputOperator* csvop);
             void addFileOutput(FileOutputOperator* fop);
 
-            inline void setOutputLimit(size_t limit) {
-                _outputLimit = limit;
+            inline void setOutputLimit(size_t topLimit, size_t bottomLimit = 0) {
+                _outputTopLimit = topLimit;
+                _outputBottomLimit = bottomLimit;
             }
 
             TransformStage* build(PhysicalPlan* plan, IBackend* backend);
@@ -134,7 +135,8 @@ namespace tuplex {
             FileFormat _outputFileFormat;
             int64_t _outputNodeID;
             int64_t _inputNodeID;
-            size_t _outputLimit;
+            size_t _outputTopLimit;
+            size_t _outputBottomLimit;
 
             LogicalOperator* _inputNode;
             std::vector<bool> _columnsToRead;
@@ -157,7 +159,7 @@ namespace tuplex {
             int64_t outputDataSetID() const;
 
             inline bool hasOutputLimit() const {
-                return _outputLimit < std::numeric_limits<size_t>::max();
+                return _outputTopLimit != std::numeric_limits<size_t>::max() && _outputBottomLimit != std::numeric_limits<size_t>::max();
             }
 
             inline char csvOutputDelimiter() const {