table.rb in red-arrow-0.16.0

- old
+ new

@@ -13,17 +13,15 @@
 # "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.
 
-require "arrow/column-containable"
-require "arrow/group"
-require "arrow/record-containable"
-
 module Arrow
   class Table
     include ColumnContainable
+    include GenericFilterable
+    include GenericTakeable
     include RecordContainable
 
     class << self
       def load(path, options={})
         TableLoader.load(path, options)
@@ -304,17 +302,17 @@
             "(given #{args.size}, expected #{expected_n_args})"
           raise ArgumentError, message
         end
       end
 
-      ranges = []
+      sliced_tables = []
       slicers.each do |slicer|
         slicer = slicer.evaluate if slicer.respond_to?(:evaluate)
         case slicer
         when Integer
           slicer += n_rows if slicer < 0
-          ranges << [slicer, n_rows - 1]
+          sliced_tables << slice_by_range(slicer, n_rows - 1)
         when Range
           original_from = from = slicer.first
           to = slicer.last
           to -= 1 if slicer.exclude_end?
           from += n_rows if from < 0
@@ -323,29 +321,25 @@
               "offset is out of range (-#{n_rows + 1},#{n_rows}): " +
               "#{original_from}"
             raise ArgumentError, message
           end
           to += n_rows if to < 0
-          ranges << [from, to]
-        when ::Array
-          boolean_array_to_slice_ranges(slicer, 0, ranges)
-        when ChunkedArray
-          offset = 0
-          slicer.each_chunk do |array|
-            boolean_array_to_slice_ranges(array, offset, ranges)
-            offset += array.length
-          end
-        when BooleanArray
-          boolean_array_to_slice_ranges(slicer, 0, ranges)
+          sliced_tables << slice_by_range(from, to)
+        when ::Array, BooleanArray, ChunkedArray
+          sliced_tables << filter(slicer)
         else
           message = "slicer must be Integer, Range, (from, to), " +
             "Arrow::ChunkedArray of Arrow::BooleanArray, " +
             "Arrow::BooleanArray or Arrow::Slicer::Condition: #{slicer.inspect}"
           raise ArgumentError, message
         end
       end
-      slice_by_ranges(ranges)
+      if sliced_tables.size > 1
+        sliced_tables[0].concatenate(sliced_tables[1..-1])
+      else
+        sliced_tables[0]
+      end
     end
 
     # TODO
     #
     # @return [Arrow::Table]
@@ -512,41 +506,11 @@
       end
       super
     end
 
     private
-    def boolean_array_to_slice_ranges(array, offset, ranges)
-      in_target = false
-      target_start = nil
-      array.each_with_index do |is_target, i|
-        if is_target
-          unless in_target
-            target_start = offset + i
-            in_target = true
-          end
-        else
-          if in_target
-            ranges << [target_start, offset + i - 1]
-            target_start = nil
-            in_target = false
-          end
-        end
-      end
-      if in_target
-        ranges << [target_start, offset + array.length - 1]
-      end
-    end
-
-    def slice_by_ranges(ranges)
-      sliced_table = []
-      ranges.each do |from, to|
-        sliced_table << slice_raw(from, to - from + 1)
-      end
-      if sliced_table.size > 1
-        sliced_table[0].concatenate(sliced_table[1..-1])
-      else
-        sliced_table[0]
-      end
+    def slice_by_range(from, to)
+      slice_raw(from, to - from + 1)
     end
 
     def ensure_raw_column(name, data)
       case data
       when Array