1# Licensed to the Apache Software Foundation (ASF) under one
2# or more contributor license agreements.  See the NOTICE file
3# distributed with this work for additional information
4# regarding copyright ownership.  The ASF licenses this file
5# to you under the Apache License, Version 2.0 (the
6# "License"); you may not use this file except in compliance
7# with the License.  You may obtain a copy of the License at
8#
9#   http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing,
12# software distributed under the License is distributed on an
13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14# KIND, either express or implied.  See the License for the
15# specific language governing permissions and limitations
16# under the License.
17
18module Arrow
19  module ColumnContainable
20    def columns
21      @columns ||= schema.n_fields.times.collect do |i|
22        Column.new(self, i)
23      end
24    end
25
26    def each_column(&block)
27      columns.each(&block)
28    end
29
30    # @overload [](name)
31    #   Find a column that has the given name.
32    #
33    #   @param name [String, Symbol] The column name to be found.
34    #   @return [Column] The found column.
35    #
36    # @overload [](index)
37    #   Find the `index`-th column.
38    #
39    #   @param index [Integer] The index to be found.
40    #   @return [Column] The found column.
41    def find_column(name_or_index)
42      case name_or_index
43      when String, Symbol
44        name = name_or_index.to_s
45        index = schema.get_field_index(name)
46        return nil if index == -1
47        Column.new(self, index)
48      when Integer
49        index = name_or_index
50        index += n_columns if index < 0
51        return nil if index < 0 or index >= n_columns
52        Column.new(self, index)
53      else
54        message = "column name or index must be String, Symbol or Integer: "
55        message << name_or_index.inspect
56        raise ArgumentError, message
57      end
58    end
59
60    # Selects columns that are selected by `selectors` and/or `block`
61    # and creates a new container only with the selected columns.
62    #
63    # @param selectors [Array<String, Symbol, Integer, Range>]
64    #   If a selector is `String`, `Symbol` or `Integer`, the selector
65    #   selects a column by {#find_column}.
66    #
67    #   If a selector is `Range`, the selector selects columns by `::Array#[]`.
68    # @yield [column] Gives a column to the block to select columns.
69    #   This uses `::Array#select`.
70    # @yieldparam column [Column] A target column.
71    # @yieldreturn [Boolean] Whether the given column is selected or not.
72    # @return [self.class] The newly created container that only has selected
73    #   columns.
74    def select_columns(*selectors, &block)
75      if selectors.empty?
76        return to_enum(__method__) unless block_given?
77        selected_columns = columns.select(&block)
78      else
79        selected_columns = []
80        selectors.each do |selector|
81          case selector
82          when Range
83            selected_columns.concat(columns[selector])
84          else
85            column = find_column(selector)
86            if column.nil?
87              case selector
88              when String, Symbol
89                message = "unknown column: #{selector.inspect}: #{inspect}"
90                raise KeyError.new(message)
91              else
92                message = "out of index (0..#{n_columns - 1}): "
93                message << "#{selector.inspect}: #{inspect}"
94                raise IndexError.new(message)
95              end
96            end
97            selected_columns << column
98          end
99        end
100        selected_columns = selected_columns.select(&block) if block_given?
101      end
102      self.class.new(selected_columns)
103    end
104
105    # @overload [](name)
106    #   Find a column that has the given name.
107    #
108    #   @param name [String, Symbol] The column name to be found.
109    #   @return [Column] The found column.
110    #   @see #find_column
111    #
112    # @overload [](index)
113    #   Find the `index`-th column.
114    #
115    #   @param index [Integer] The index to be found.
116    #   @return [Column] The found column.
117    #   @see #find_column
118    #
119    # @overload [](range)
120    #   Selects columns that are in `range` and creates a new container
121    #   only with the selected columns.
122    #
123    #   @param range [Range] The range to be selected.
124    #   @return [self.class] The newly created container that only has selected
125    #     columns.
126    #   @see #select_columns
127    #
128    # @overload [](selectors)
129    #   Selects columns that are selected by `selectors` and creates a
130    #   new container only with the selected columns.
131    #
132    #   @param selectors [Array] The selectors that are used to select columns.
133    #   @return [self.class] The newly created container that only has selected
134    #     columns.
135    #   @see #select_columns
136    def [](selector)
137      case selector
138      when ::Array
139        select_columns(*selector)
140      when Range
141        select_columns(selector)
142      else
143        find_column(selector)
144      end
145    end
146  end
147end
148