lib/arrow/table.rb in red-arrow-0.8.2 vs lib/arrow/table.rb in red-arrow-0.10.0

@@ -1,18 +1,21 @@
-# Copyright 2017-2018 Kouhei Sutou <[e-mail]>
+# 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
 #
-# Licensed 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
 #
-#     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.
+# 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.
 
 require "arrow/group"
 require "arrow/record-containable"
 
 module Arrow
@@ -68,11 +71,34 @@
     alias_method :[], :find_column
 
     # TODO
     #
     # @return [Arrow::Table]
-    def slice(*slicers)
+    def slice(*args)
+      slicers = []
+      expected_n_args = nil
+      case args.size
+      when 0
+        expected_n_args = "1..2" unless block_given?
+      when 1
+        slicers << args[0]
+      when 2
+        from, to = args
+        slicers << (from...(from + to))
+      else
+        if block_given?
+          expected_n_args = "0..2"
+        else
+          expected_n_args = "1..2"
+        end
+      end
+      if expected_n_args
+        message = "wrong number of arguments " +
+          "(given #{args.size}, expected #{expected_n_args})"
+        raise ArgumentError, message
+      end
+
       if block_given?
         block_slicer = yield(Slicer.new(self))
         case block_slicer
         when nil
           # Ignore
@@ -80,10 +106,11 @@
           slicers.concat(block_slicer)
         else
           slicers << block_slicer
         end
       end
+
       ranges = []
       slicers.each do |slicer|
         slicer = slicer.evaluate if slicer.respond_to?(:evaluate)
         case slicer
         when Integer
@@ -95,24 +122,21 @@
           to -= 1 if slicer.exclude_end?
           from += n_rows if from < 0
           to += n_rows if to < 0
           ranges << [from, to]
         when ::Array
-          from = slicer[0]
-          from += n_rows if from < 0
-          to = from + slicer[1] - 1
-          ranges << [from, to]
+          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)
         else
-          message = "slicer must be Integer, Range, [from, to], " +
+          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
@@ -222,14 +246,20 @@
         selected_columns = selected_columns.select(&block) if block_given?
       end
       self.class.new(selected_columns)
     end
 
+    # Experimental
     def group(*keys)
       Group.new(self, keys)
     end
 
+    # Experimental
+    def window(size: nil)
+      RollingWindow.new(self, size)
+    end
+
     def save(path, options={})
       saver = TableSaver.new(self, path, options)
       saver.save
     end
 
@@ -289,9 +319,10 @@
       if in_target
         ranges << [target_start, offset + array.length - 1]
       end
     end
 
+    # TODO: Almost codes should be implemented in Apache Arrow C++.
     def slice_by_ranges(ranges)
       sliced_columns = columns.collect do |column|
         chunks = []
         arrays = column.data.each_chunk.to_a
         offset = 0