Skip to content

Commit 0249a2e

Browse files
authored
GH-50225: [Ruby] Move merge implementation to ColumnContainable (#50226)
### Rationale for this change `Arrow::Table#merge` and `Arrow::RecordBatch#merge` had similar implementations, resulting in duplicated merge logic. ### What changes are included in this PR? This PR moves the common merge implementation to `Arrow::ColumnContainable`. Container-specific behavior such as column conversion and creating the merged container remains implemented in `Arrow::Table` and `Arrow::RecordBatch`. ### Are these changes tested? Yes. The existing `RecordBatch` and `Table` test suites were run and passed successfully. ### Are there any user-facing changes? No. * GitHub Issue: #50225 Authored-by: Aaditya Srinivasan <aadityasri03@gmail.com> Signed-off-by: Sutou Kouhei <kou@clear-code.com>
1 parent d6bba6b commit 0249a2e

3 files changed

Lines changed: 71 additions & 112 deletions

File tree

ruby/red-arrow/lib/arrow/column-containable.rb

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,5 +152,70 @@ def [](selector)
152152
def column_names
153153
@column_names ||= columns.collect(&:name)
154154
end
155+
156+
# Merges columns from the given container or Hash and creates
157+
# a new container.
158+
#
159+
# @param other [Hash, self]
160+
# The columns to be merged.
161+
#
162+
# @return [self]
163+
def merge(other)
164+
added_columns = {}
165+
removed_columns = {}
166+
167+
case other
168+
when Hash
169+
other.each do |name, value|
170+
name = name.to_s
171+
if value
172+
added_columns[name] = ensure_raw_column(name, value)
173+
else
174+
removed_columns[name] = true
175+
end
176+
end
177+
when self.class
178+
other.columns.each do |column|
179+
name = column.name
180+
added_columns[name] = ensure_raw_column(name, column)
181+
end
182+
else
183+
message = "merge target must be Hash or #{self.class}: " +
184+
"<#{other.inspect}>: #{inspect}"
185+
raise ArgumentError, message
186+
end
187+
188+
new_columns = []
189+
190+
columns.each do |column|
191+
column_name = column.name
192+
new_column = added_columns.delete(column_name)
193+
194+
if new_column
195+
new_columns << new_column
196+
next
197+
end
198+
199+
next if removed_columns.key?(column_name)
200+
201+
new_columns << ensure_raw_column(column_name, column)
202+
end
203+
204+
added_columns.each_value do |new_column|
205+
new_columns << new_column
206+
end
207+
208+
new_fields = []
209+
new_arrays = []
210+
211+
new_columns.each do |new_column|
212+
new_fields << new_column[:field]
213+
new_arrays << new_column[:data]
214+
end
215+
216+
merged = self.class.new(new_fields, new_arrays)
217+
share_input(merged)
218+
merged
219+
end
155220
end
156221
end

ruby/red-arrow/lib/arrow/record-batch.rb

Lines changed: 6 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,12 @@ def new(*args)
3737
super(schema, n_rows, values)
3838
when 2
3939
schema, data = args
40-
RecordBatchBuilder.build(schema, data)
40+
schema = Schema.new(schema) unless schema.is_a?(Schema)
41+
if !data.empty? and data.all? {|array| array.is_a?(Arrow::Array)}
42+
super(schema, data[0].size, data)
43+
else
44+
RecordBatchBuilder.build(schema, data)
45+
end
4146
when 3
4247
super
4348
else
@@ -63,63 +68,6 @@ def to_table
6368
table
6469
end
6570

66-
def merge(other)
67-
added_columns = {}
68-
removed_columns = {}
69-
70-
case other
71-
when Hash
72-
other.each do |name, value|
73-
name = name.to_s
74-
if value
75-
added_columns[name] = ensure_raw_column(name, value)
76-
else
77-
removed_columns[name] = true
78-
end
79-
end
80-
when RecordBatch
81-
other.columns.each do |column|
82-
name = column.name
83-
added_columns[name] = ensure_raw_column(name, column)
84-
end
85-
else
86-
message = "merge target must be Hash or Arrow::RecordBatch: " +
87-
"<#{other.inspect}>: #{inspect}"
88-
raise ArgumentError, message
89-
end
90-
91-
new_columns = []
92-
columns.each do |column|
93-
column_name = column.name
94-
new_column = added_columns.delete(column_name)
95-
if new_column
96-
new_columns << new_column
97-
next
98-
end
99-
next if removed_columns.key?(column_name)
100-
new_columns << ensure_raw_column(column_name, column)
101-
end
102-
103-
added_columns.each_value do |new_column|
104-
new_columns << new_column
105-
end
106-
107-
new_fields = []
108-
new_arrays = []
109-
new_columns.each do |new_column|
110-
new_fields << new_column[:field]
111-
new_arrays << new_column[:data]
112-
end
113-
114-
record_batch = self.class.new(
115-
Schema.new(new_fields),
116-
n_rows,
117-
new_arrays,
118-
)
119-
share_input(record_batch)
120-
record_batch
121-
end
122-
12371
def respond_to_missing?(name, include_private)
12472
return true if find_column(name)
12573
super

ruby/red-arrow/lib/arrow/table.rb

Lines changed: 0 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -354,60 +354,6 @@ def slice(*args)
354354
sliced_table
355355
end
356356

357-
# TODO
358-
#
359-
# @return [Arrow::Table]
360-
def merge(other)
361-
added_columns = {}
362-
removed_columns = {}
363-
364-
case other
365-
when Hash
366-
other.each do |name, value|
367-
name = name.to_s
368-
if value
369-
added_columns[name] = ensure_raw_column(name, value)
370-
else
371-
removed_columns[name] = true
372-
end
373-
end
374-
when Table
375-
added_columns = {}
376-
other.columns.each do |column|
377-
name = column.name
378-
added_columns[name] = ensure_raw_column(name, column)
379-
end
380-
else
381-
message = "merge target must be Hash or Arrow::Table: " +
382-
"<#{other.inspect}>: #{inspect}"
383-
raise ArgumentError, message
384-
end
385-
386-
new_columns = []
387-
columns.each do |column|
388-
column_name = column.name
389-
new_column = added_columns.delete(column_name)
390-
if new_column
391-
new_columns << new_column
392-
next
393-
end
394-
next if removed_columns.key?(column_name)
395-
new_columns << ensure_raw_column(column_name, column)
396-
end
397-
added_columns.each do |name, new_column|
398-
new_columns << new_column
399-
end
400-
new_fields = []
401-
new_arrays = []
402-
new_columns.each do |new_column|
403-
new_fields << new_column[:field]
404-
new_arrays << new_column[:data]
405-
end
406-
table = self.class.new(new_fields, new_arrays)
407-
share_input(table)
408-
table
409-
end
410-
411357
alias_method :remove_column_raw, :remove_column
412358
def remove_column(name_or_index)
413359
case name_or_index

0 commit comments

Comments
 (0)