]>
Commit | Line | Data |
---|---|---|
1d09f67e TL |
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 | ||
18 | module 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 | |
147 | end |