Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Commit

Permalink
Simpler IPC code (#939)
Browse files Browse the repository at this point in the history
  • Loading branch information
jorgecarleitao committed Apr 12, 2022
1 parent 0d93c58 commit 60130f4
Show file tree
Hide file tree
Showing 3 changed files with 41 additions and 107 deletions.
15 changes: 8 additions & 7 deletions src/io/ipc/write/common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,20 +45,21 @@ fn encode_dictionary(
let dict_id = field.dictionary_id
.ok_or_else(|| ArrowError::InvalidArgumentError("Dictionaries must have an associated id".to_string()))?;

let values = array.as_any().downcast_ref::<DictionaryArray<$T>>().unwrap().values();
let emit = dictionary_tracker.insert(dict_id, array)?;

let array = array.as_any().downcast_ref::<DictionaryArray<$T>>().unwrap();
let values = array.values();
encode_dictionary(field,
values,
options,
dictionary_tracker,
encoded_dictionaries
)?;

let emit = dictionary_tracker.insert(dict_id, array)?;

if emit {
encoded_dictionaries.push(dictionary_batch_to_bytes(
encoded_dictionaries.push(dictionary_batch_to_bytes::<$T>(
dict_id,
array.as_ref(),
array,
options,
is_native_little_endian(),
));
Expand Down Expand Up @@ -257,9 +258,9 @@ fn columns_to_bytes(columns: &Chunk<Arc<dyn Array>>, options: &WriteOptions) ->

/// Write dictionary values into two sets of bytes, one for the header (ipc::Schema::Message) and the
/// other for the data
fn dictionary_batch_to_bytes(
fn dictionary_batch_to_bytes<K: DictionaryKey>(
dict_id: i64,
array: &dyn Array,
array: &DictionaryArray<K>,
options: &WriteOptions,
is_little_endian: bool,
) -> EncodedData {
Expand Down
3 changes: 2 additions & 1 deletion src/io/ipc/write/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,8 @@ mod writer;

pub use common::{Compression, Record, WriteOptions};
pub use schema::schema_to_bytes;
pub use serialize::{write, write_dictionary};
pub use serialize::write;
pub(self) use serialize::write_dictionary;
pub use stream::StreamWriter;
pub use writer::FileWriter;

Expand Down
Loading

0 comments on commit 60130f4

Please sign in to comment.