This is an automated email from the ASF dual-hosted git repository.
kou pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new 237705bf17 GH-34440: [Ruby] Add support for
`RecordBatch{File,Stream}Reader#each` without block (#34441)
237705bf17 is described below
commit 237705bf17486cfc35ab7d1ddfe59dd60f042ab8
Author: Sutou Kouhei <[email protected]>
AuthorDate: Sun Mar 5 07:44:37 2023 +0900
GH-34440: [Ruby] Add support for `RecordBatch{File,Stream}Reader#each`
without block (#34441)
### Rationale for this change
`#each` must support this.
### What changes are included in this PR?
`to_enum` and tests for this.
### Are these changes tested?
Yes.
### Are there any user-facing changes?
Yes.
* Closes: #34440
Lead-authored-by: Sutou Kouhei <[email protected]>
Co-authored-by: Sutou Kouhei <[email protected]>
Co-authored-by: Hirokazu SUZUKI <[email protected]>
Signed-off-by: Sutou Kouhei <[email protected]>
---
.../lib/arrow/record-batch-file-reader.rb | 2 ++
.../lib/arrow/record-batch-stream-reader.rb | 2 ++
.../test/test-record-batch-file-reader.rb | 21 +++++++++++++++++++++
...eader.rb => test-record-batch-stream-reader.rb} | 22 ++++++++++++++++++----
4 files changed, 43 insertions(+), 4 deletions(-)
diff --git a/ruby/red-arrow/lib/arrow/record-batch-file-reader.rb
b/ruby/red-arrow/lib/arrow/record-batch-file-reader.rb
index 86a757e320..5f684a6cfc 100644
--- a/ruby/red-arrow/lib/arrow/record-batch-file-reader.rb
+++ b/ruby/red-arrow/lib/arrow/record-batch-file-reader.rb
@@ -20,6 +20,8 @@ module Arrow
include Enumerable
def each
+ return to_enum(__method__) {n_record_batches} unless block_given?
+
n_record_batches.times do |i|
yield(get_record_batch(i))
end
diff --git a/ruby/red-arrow/lib/arrow/record-batch-stream-reader.rb
b/ruby/red-arrow/lib/arrow/record-batch-stream-reader.rb
index fa15c80009..44341a6745 100644
--- a/ruby/red-arrow/lib/arrow/record-batch-stream-reader.rb
+++ b/ruby/red-arrow/lib/arrow/record-batch-stream-reader.rb
@@ -20,6 +20,8 @@ module Arrow
include Enumerable
def each
+ return to_enum(__method__) unless block_given?
+
loop do
record_batch = next_record_batch
break if record_batch.nil?
diff --git a/ruby/red-arrow/test/test-record-batch-file-reader.rb
b/ruby/red-arrow/test/test-record-batch-file-reader.rb
index 57b02abf9c..7f7fb79c08 100644
--- a/ruby/red-arrow/test/test-record-batch-file-reader.rb
+++ b/ruby/red-arrow/test/test-record-batch-file-reader.rb
@@ -112,4 +112,25 @@ class RecordBatchFileReaderTest < Test::Unit::TestCase
end
end
end
+
+ sub_test_case("#each") do
+ test("without block") do
+ buffer = Arrow::ResizableBuffer.new(1024)
+ Arrow::Table.new(number: [1, 2, 3]).save(buffer)
+ Arrow::BufferInputStream.open(buffer) do |input|
+ reader = Arrow::RecordBatchFileReader.new(input)
+ each = reader.each
+ assert_equal({
+ size: 1,
+ to_a: [
+ Arrow::RecordBatch.new(number: [1, 2, 3]),
+ ],
+ },
+ {
+ size: each.size,
+ to_a: each.to_a,
+ })
+ end
+ end
+ end
end
diff --git a/ruby/red-arrow/test/test-record-batch-file-reader.rb
b/ruby/red-arrow/test/test-record-batch-stream-reader.rb
similarity index 85%
copy from ruby/red-arrow/test/test-record-batch-file-reader.rb
copy to ruby/red-arrow/test/test-record-batch-stream-reader.rb
index 57b02abf9c..6db2dd4925 100644
--- a/ruby/red-arrow/test/test-record-batch-file-reader.rb
+++ b/ruby/red-arrow/test/test-record-batch-stream-reader.rb
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
-class RecordBatchFileReaderTest < Test::Unit::TestCase
+class RecordBatchStreamReaderTest < Test::Unit::TestCase
test("write/read") do
fields = [
Arrow::Field.new("uint8", :uint8),
@@ -31,9 +31,9 @@ class RecordBatchFileReaderTest < Test::Unit::TestCase
]
schema = Arrow::Schema.new(fields)
- tempfile = Tempfile.new(["batch", ".arrow"])
+ tempfile = Tempfile.new(["batch", ".arrows"])
Arrow::FileOutputStream.open(tempfile.path, false) do |output|
- Arrow::RecordBatchFileWriter.open(output, schema) do |writer|
+ Arrow::RecordBatchStreamWriter.open(output, schema) do |writer|
uints = [1, 2, 4, 8]
ints = [1, -2, 4, -8]
floats = [1.1, -2.2, 4.4, -8.8]
@@ -56,7 +56,7 @@ class RecordBatchFileReaderTest < Test::Unit::TestCase
end
Arrow::MemoryMappedInputStream.open(tempfile.path) do |input|
- reader = Arrow::RecordBatchFileReader.new(input)
+ reader = Arrow::RecordBatchStreamReader.new(input)
reader.each do |record_batch|
assert_equal([
{
@@ -112,4 +112,18 @@ class RecordBatchFileReaderTest < Test::Unit::TestCase
end
end
end
+
+ sub_test_case("#each") do
+ test("without block") do
+ buffer = Arrow::ResizableBuffer.new(1024)
+ Arrow::Table.new(number: [1, 2, 3]).save(buffer, format: :arrows)
+ Arrow::BufferInputStream.open(buffer) do |input|
+ reader = Arrow::RecordBatchStreamReader.new(input)
+ assert_equal([
+ Arrow::RecordBatch.new(number: [1, 2, 3]),
+ ],
+ reader.each.to_a)
+ end
+ end
+ end
end
