Skip to content

Commit

Permalink
Remove deprecated flight_data_from_arrow_batch (#6823)
Browse files Browse the repository at this point in the history
`flight_data_from_arrow_batch` was deprecated since v 30.0.0.
The only remaining usage was in tests, so convert it to a test helper.
  • Loading branch information
findepi authored Dec 3, 2024
1 parent b3f9355 commit e0906ae
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 27 deletions.
23 changes: 20 additions & 3 deletions arrow-flight/src/encode.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1586,12 +1586,29 @@ mod tests {
hydrate_dictionaries(&batch, batch.schema()).expect("failed to optimize");
}

pub fn make_flight_data(
fn make_flight_data(
batch: &RecordBatch,
options: &IpcWriteOptions,
) -> (Vec<FlightData>, FlightData) {
#[allow(deprecated)]
crate::utils::flight_data_from_arrow_batch(batch, options)
flight_data_from_arrow_batch(batch, options)
}

fn flight_data_from_arrow_batch(
batch: &RecordBatch,
options: &IpcWriteOptions,
) -> (Vec<FlightData>, FlightData) {
let data_gen = IpcDataGenerator::default();
let mut dictionary_tracker =
DictionaryTracker::new_with_preserve_dict_id(false, options.preserve_dict_id());

let (encoded_dictionaries, encoded_batch) = data_gen
.encoded_batch(batch, &mut dictionary_tracker, options)
.expect("DictionaryTracker configured above to not error on replacement");

let flight_dictionaries = encoded_dictionaries.into_iter().map(Into::into).collect();
let flight_batch = encoded_batch.into();

(flight_dictionaries, flight_batch)
}

#[test]
Expand Down
24 changes: 0 additions & 24 deletions arrow-flight/src/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,30 +27,6 @@ use arrow_ipc::convert::fb_to_schema;
use arrow_ipc::{reader, root_as_message, writer, writer::IpcWriteOptions};
use arrow_schema::{ArrowError, Schema, SchemaRef};

/// Convert a `RecordBatch` to a vector of `FlightData` representing the bytes of the dictionaries
/// and a `FlightData` representing the bytes of the batch's values
#[deprecated(
since = "30.0.0",
note = "Use IpcDataGenerator directly with DictionaryTracker to avoid re-sending dictionaries"
)]
pub fn flight_data_from_arrow_batch(
batch: &RecordBatch,
options: &IpcWriteOptions,
) -> (Vec<FlightData>, FlightData) {
let data_gen = writer::IpcDataGenerator::default();
let mut dictionary_tracker =
writer::DictionaryTracker::new_with_preserve_dict_id(false, options.preserve_dict_id());

let (encoded_dictionaries, encoded_batch) = data_gen
.encoded_batch(batch, &mut dictionary_tracker, options)
.expect("DictionaryTracker configured above to not error on replacement");

let flight_dictionaries = encoded_dictionaries.into_iter().map(Into::into).collect();
let flight_batch = encoded_batch.into();

(flight_dictionaries, flight_batch)
}

/// Convert a slice of wire protocol `FlightData`s into a vector of `RecordBatch`es
pub fn flight_data_to_batches(flight_data: &[FlightData]) -> Result<Vec<RecordBatch>, ArrowError> {
let schema = flight_data.first().ok_or_else(|| {
Expand Down

0 comments on commit e0906ae

Please sign in to comment.