red-arrow 0.15.1 → 1.0.1

data/lib/arrow/raw-table-converter.rb

@@ -0,0 +1,47 @@
+# 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
+#
+#   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.
+
+module Arrow
+  class RawTableConverter
+    attr_reader :n_rows
+    attr_reader :schema
+    attr_reader :values
+    def initialize(raw_table)
+      @raw_table = raw_table
+      convert
+    end
+
+    private
+    def convert
+      if @raw_table.is_a?(::Array) and @raw_table[0].is_a?(Column)
+        fields = @raw_table.collect(&:field)
+        @schema = Schema.new(fields)
+        @values = @raw_table.collect(&:data)
+      else
+        fields = []
+        @values = []
+        @raw_table.each do |name, array|
+          array = ArrayBuilder.build(array) if array.is_a?(::Array)
+          fields << Field.new(name.to_s, array.value_data_type)
+          @values << array
+        end
+        @schema = Schema.new(fields)
+      end
+      @n_rows = @values[0].length
+    end
+  end
+end

data/lib/arrow/record-batch-iterator.rb

@@ -0,0 +1,22 @@
+# 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
+#
+#   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.
+
+module Arrow
+  class RecordBatchIterator
+    alias_method :to_a, :to_list
+  end
+end

data/lib/arrow/record-batch.rb

@@ -15,8 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-require "arrow/column-containable"
-require "arrow/record-containable"
+require "arrow/raw-table-converter"
 
 module Arrow
   class RecordBatch
@@ -28,13 +27,19 @@ module Arrow
       def new(*args)
         n_args = args.size
         case n_args
+        when 1
+          raw_table_converter = RawTableConverter.new(args[0])
+          n_rows = raw_table_converter.n_rows
+          schema = raw_table_converter.schema
+          values = raw_table_converter.values
+          super(schema, n_rows, values)
         when 2
           schema, data = args
           RecordBatchBuilder.build(schema, data)
         when 3
           super
         else
-          message = "wrong number of arguments (given #{n_args}, expected 2..3)"
+          message = "wrong number of arguments (given #{n_args}, expected 1..3)"
           raise ArgumentError, message
         end
       end

data/lib/arrow/schema.rb

@@ -15,8 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 
-require "arrow/field-containable"
-
 module Arrow
   class Schema
     include FieldContainable
@@ -93,5 +91,10 @@ module Arrow
     end
 
     alias_method :[], :find_field
+
+    alias_method :to_s_raw, :to_s
+    def to_s(show_metadata: false)
+      to_string_metadata(show_metadata)
+    end
   end
 end

data/lib/arrow/struct-array-builder.rb

@@ -32,7 +32,7 @@ module Arrow
       case index_or_name
       when String, Symbol
         name = index_or_name
-        (@name_to_builder ||= build_name_to_builder)[name.to_s]
+        cached_name_to_builder[name.to_s]
       else
         index = index_or_name
         cached_field_builders[index]
@@ -70,13 +70,18 @@ module Arrow
           append_null
         when ::Array
           append_value_raw
-          value.each_with_index do |sub_value, i|
-            self[i].append(sub_value)
+          cached_field_builders.zip(value) do |builder, sub_value|
+            builder.append(sub_value)
           end
         when Hash
           append_value_raw
+          local_name_to_builder = cached_name_to_builder.dup
           value.each do |name, sub_value|
-            self[name].append(sub_value)
+            builder = local_name_to_builder.delete(name.to_s)
+            builder.append(sub_value)
+          end
+          local_name_to_builder.each do |_, builder|
+            builder.append_null
           end
         else
           message =
@@ -108,9 +113,6 @@ module Arrow
     alias_method :append_null_raw, :append_null
     def append_null
       append_null_raw
-      cached_field_builders.each do |builder|
-        builder.append_null
-      end
     end
 
     # @since 0.12.0
@@ -136,5 +138,9 @@ module Arrow
       end
       name_to_builder
     end
+
+    def cached_name_to_builder
+      @name_to_builder ||= build_name_to_builder
+    end
   end
 end

data/lib/arrow/struct-data-type.rb

@@ -15,8 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 
-require "arrow/field-containable"
-
 module Arrow
   class StructDataType
     include FieldContainable

data/lib/arrow/table-loader.rb

@@ -41,6 +41,8 @@ module Arrow
             available_formats << match_data.post_match
           end
         end
+        deprecated_formats = ["batch", "stream"]
+        available_formats -= deprecated_formats
         message = "Arrow::Table load format must be one of ["
         message << available_formats.join(", ")
         message << "]: #{format.inspect}"
@@ -119,18 +121,30 @@ module Arrow
       load_raw(input, reader)
     end
 
-    def load_as_batch
+    # @since 1.0.0
+    def load_as_arrow_file
       input = open_input_stream
       reader = RecordBatchFileReader.new(input)
       load_raw(input, reader)
     end
 
-    def load_as_stream
+    # @deprecated Use `format: :arrow_file` instead.
+    def load_as_batch
+      load_as_arrow_file
+    end
+
+    # @since 1.0.0
+    def load_as_arrow_streaming
       input = open_input_stream
       reader = RecordBatchStreamReader.new(input)
       load_raw(input, reader)
     end
 
+    # @deprecated Use `format: :arrow_streaming` instead.
+    def load_as_stream
+      load_as_arrow_streaming
+    end
+
     if Arrow.const_defined?(:ORCFileReader)
       def load_as_orc
         input = open_input_stream
@@ -143,16 +157,25 @@ module Arrow
       end
     end
 
-    def load_as_csv
-      options = @options.dup
+    def csv_load(options)
       options.delete(:format)
       if @input.is_a?(Buffer)
-        CSVLoader.load(@input.data.to_s, options)
+        CSVLoader.load(@input.data.to_s, **options)
       else
-        CSVLoader.load(Pathname.new(@input), options)
+        CSVLoader.load(Pathname.new(@input), **options)
       end
     end
 
+    def load_as_csv
+      csv_load(@options.dup)
+    end
+
+    def load_as_tsv
+      options = @options.dup
+      options[:delimiter] = "\t"
+      csv_load(options.dup)
+    end
+
     def load_as_feather
       input = open_input_stream
       reader = FeatherFileReader.new(input)

data/lib/arrow/table-saver.rb

@@ -42,6 +42,8 @@ module Arrow
             available_formats << match_data.post_match
           end
         end
+        deprecated_formats = ["batch", "stream"]
+        available_formats -= deprecated_formats
         message = "Arrow::Table save format must be one of ["
         message << available_formats.join(", ")
         message << "]: #{format.inspect}"
@@ -110,35 +112,57 @@ module Arrow
     end
 
     def save_as_arrow
-      save_as_batch
+      save_as_arrow_file
     end
 
-    def save_as_batch
+    # @since 1.0.0
+    def save_as_arrow_file
       save_raw(RecordBatchFileWriter)
     end
 
-    def save_as_stream
+    # @deprecated Use `format: :arrow_batch` instead.
+    def save_as_batch
+      save_as_arrow_file
+    end
+
+    # @since 1.0.0
+    def save_as_arrow_streaming
       save_raw(RecordBatchStreamWriter)
     end
 
-    def save_as_csv
+    # @deprecated Use `format: :arrow_streaming` instead.
+    def save_as_stream
+      save_as_arrow_streaming
+    end
+
+    def csv_save(**options)
       open_output_stream do |output|
-        csv = CSV.new(output)
+        csv = CSV.new(output, **options)
         names = @table.schema.fields.collect(&:name)
         csv << names
-        @table.each_record(reuse_record: true) do |record|
-          csv << names.collect do |name|
-            record[name]
-          end
+        @table.raw_records.each do |record|
+          csv << record
         end
       end
     end
 
+    def save_as_csv
+      csv_save
+    end
+
+    def save_as_tsv
+      csv_save(col_sep: "\t")
+    end
+
     def save_as_feather
-      open_output_stream do |output|
-        FeatherFileWriter.open(output) do |writer|
-          writer.write(@table)
-        end
+      properties = FeatherWriteProperties.new
+      properties.class.properties.each do |name|
+        value = @options[name.to_sym]
+        next if value.nil?
+        properties.__send__("#{name}=", value)
+      end
+      open_raw_output_stream do |output|
+        @table.write_as_feather(output, properties)
       end
     end
   end

data/lib/arrow/table.rb

@@ -15,13 +15,13 @@
 # specific language governing permissions and limitations
 # under the License.
 
-require "arrow/column-containable"
-require "arrow/group"
-require "arrow/record-containable"
+require "arrow/raw-table-converter"
 
 module Arrow
   class Table
     include ColumnContainable
+    include GenericFilterable
+    include GenericTakeable
     include RecordContainable
 
     class << self
@@ -83,14 +83,6 @@ module Arrow
     #     `Array`.
     #
     #   @example Create a table from column name and values
-    #     count_chunks = [
-    #       Arrow::UInt32Array.new([0, 2]),
-    #       Arrow::UInt32Array.new([nil, 4]),
-    #     ]
-    #     visible_chunks = [
-    #       Arrow::BooleanArray.new([true]),
-    #       Arrow::BooleanArray.new([nil, nil, false]),
-    #     ]
     #     Arrow::Table.new("count" => [0, 2, nil, 4],
     #                      "visible" => [true, nil, nil, false])
     #
@@ -171,22 +163,9 @@ module Arrow
       n_args = args.size
       case n_args
       when 1
-        if args[0][0].is_a?(Column)
-          columns = args[0]
-          fields = columns.collect(&:field)
-          values = columns.collect(&:data)
-          schema = Schema.new(fields)
-        else
-          raw_table = args[0]
-          fields = []
-          values = []
-          raw_table.each do |name, array|
-            array = ArrayBuilder.build(array) if array.is_a?(::Array)
-            fields << Field.new(name.to_s, array.value_data_type)
-            values << array
-          end
-          schema = Schema.new(fields)
-        end
+        raw_table_converter = RawTableConverter.new(args[0])
+        schema = raw_table_converter.schema
+        values = raw_table_converter.values
       when 2
         schema = args[0]
         schema = Schema.new(schema) unless schema.is_a?(Schema)
@@ -306,13 +285,15 @@ module Arrow
         end
       end
 
-      ranges = []
+      filter_options = Arrow::FilterOptions.new
+      filter_options.null_selection_behavior = :emit_null
+      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
@@ -325,17 +306,9 @@ module Arrow
             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, filter_options)
         else
           message = "slicer must be Integer, Range, (from, to), " +
             "Arrow::ChunkedArray of Arrow::BooleanArray, " +
@@ -343,7 +316,11 @@ module Arrow
           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
@@ -514,38 +491,8 @@ module Arrow
     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)