Skip to content

Commit ceb0659

Browse files
author
Devdutt Shenoi
committed
test: no manual_write and fix read order
1 parent 01abc11 commit ceb0659

File tree

1 file changed

+18
-52
lines changed

1 file changed

+18
-52
lines changed

src/parseable/staging/reader.rs

Lines changed: 18 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -133,10 +133,7 @@ mod tests {
133133
use arrow_array::{
134134
cast::AsArray, types::Int64Type, Array, Float64Array, Int64Array, RecordBatch, StringArray,
135135
};
136-
use arrow_ipc::{
137-
reader::FileReader,
138-
writer::{write_message, DictionaryTracker, FileWriter, IpcDataGenerator, IpcWriteOptions},
139-
};
136+
use arrow_ipc::{reader::FileReader, writer::FileWriter};
140137
use temp_dir::TempDir;
141138

142139
fn rb(rows: usize) -> RecordBatch {
@@ -197,64 +194,33 @@ mod tests {
197194
let reader = File::open(path).unwrap();
198195
let mut reader = FileReader::try_new_buffered(reader, None).unwrap();
199196
let rb = reader.next().unwrap().unwrap();
200-
assert_eq!(rb.num_rows(), 3);
197+
assert_eq!(rb.num_rows(), 1);
201198
let col1_val: Vec<i64> = rb
202199
.column(0)
203200
.as_primitive::<Int64Type>()
204201
.iter()
205202
.flatten()
206203
.collect();
207-
assert_eq!(col1_val, vec![0, 1, 2]);
204+
assert_eq!(col1_val, vec![0]);
208205

209206
let rb = reader.next().unwrap().unwrap();
210207
assert_eq!(rb.num_rows(), 2);
208+
let col1_val: Vec<i64> = rb
209+
.column(0)
210+
.as_primitive::<Int64Type>()
211+
.iter()
212+
.flatten()
213+
.collect();
214+
assert_eq!(col1_val, vec![0, 1]);
211215

212216
let rb = reader.next().unwrap().unwrap();
213-
assert_eq!(rb.num_rows(), 1);
214-
}
215-
216-
#[test]
217-
fn manual_write() {
218-
let temp_dir = TempDir::new().unwrap();
219-
let path = temp_dir.path().join("test.arrows");
220-
let error_on_replacement = true;
221-
let options = IpcWriteOptions::default();
222-
let mut dictionary_tracker = DictionaryTracker::new(error_on_replacement);
223-
let data_gen = IpcDataGenerator {};
224-
let mut file = File::create(&path).unwrap();
225-
let rb1 = rb(1);
226-
227-
let schema = data_gen.schema_to_bytes_with_dictionary_tracker(
228-
&rb1.schema(),
229-
&mut dictionary_tracker,
230-
&options,
231-
);
232-
write_message(&mut file, schema, &options).unwrap();
233-
234-
for i in (1..=3).cycle().skip(1).take(10000) {
235-
let (_, encoded_message) = data_gen
236-
.encoded_batch(&rb(i), &mut dictionary_tracker, &options)
237-
.unwrap();
238-
write_message(&mut file, encoded_message, &options).unwrap();
239-
}
240-
241-
let schema = data_gen.schema_to_bytes_with_dictionary_tracker(
242-
&rb1.schema(),
243-
&mut dictionary_tracker,
244-
&options,
245-
);
246-
write_message(&mut file, schema, &options).unwrap();
247-
248-
let reader = File::open(path).unwrap();
249-
let reader = FileReader::try_new_buffered(reader, None).unwrap();
250-
251-
let mut sum = 0;
252-
for rb in reader {
253-
let rb = rb.unwrap();
254-
sum += 1;
255-
assert!(rb.num_rows() > 0);
256-
}
257-
258-
assert_eq!(sum, 10000);
217+
assert_eq!(rb.num_rows(), 3);
218+
let col1_val: Vec<i64> = rb
219+
.column(0)
220+
.as_primitive::<Int64Type>()
221+
.iter()
222+
.flatten()
223+
.collect();
224+
assert_eq!(col1_val, vec![0, 1, 2]);
259225
}
260226
}

0 commit comments

Comments
 (0)