You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ag...@apache.org on 2024/08/12 19:59:18 UTC
(datafusion-comet) branch main updated: Minor: ScanExec code cleanup and additional documentation (#804)
This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 68f41f65 Minor: ScanExec code cleanup and additional documentation (#804)
68f41f65 is described below
commit 68f41f652b8d80530c5c0c20e177c9b84e8d2474
Author: Andy Grove <ag...@apache.org>
AuthorDate: Mon Aug 12 13:59:14 2024 -0600
Minor: ScanExec code cleanup and additional documentation (#804)
* Minor code cleanup and adding additional documentation to ScanExec
* use for loop instead of mutating an index in a while loop
* stop timer on early return
* re-use CastOptions across batches
* add LargeBinary handling to unpack_dictionary_type
---
native/core/src/execution/datafusion/planner.rs | 5 ++-
native/core/src/execution/operators/scan.rs | 59 ++++++++++++++++---------
2 files changed, 40 insertions(+), 24 deletions(-)
diff --git a/native/core/src/execution/datafusion/planner.rs b/native/core/src/execution/datafusion/planner.rs
index e1d65a4b..c3aad61e 100644
--- a/native/core/src/execution/datafusion/planner.rs
+++ b/native/core/src/execution/datafusion/planner.rs
@@ -867,7 +867,7 @@ impl PhysicalPlanner {
))
}
OpStruct::Scan(scan) => {
- let fields = scan.fields.iter().map(to_arrow_datatype).collect_vec();
+ let data_types = scan.fields.iter().map(to_arrow_datatype).collect_vec();
// If it is not test execution context for unit test, we should have at least one
// input source
@@ -887,7 +887,8 @@ impl PhysicalPlanner {
};
// The `ScanExec` operator will take actual arrays from Spark during execution
- let scan = ScanExec::new(self.exec_context_id, input_source, &scan.source, fields)?;
+ let scan =
+ ScanExec::new(self.exec_context_id, input_source, &scan.source, data_types)?;
Ok((vec![scan.clone()], Arc::new(scan)))
}
OpStruct::ShuffleWriter(writer) => {
diff --git a/native/core/src/execution/operators/scan.rs b/native/core/src/execution/operators/scan.rs
index ae19b368..0d5b3393 100644
--- a/native/core/src/execution/operators/scan.rs
+++ b/native/core/src/execution/operators/scan.rs
@@ -50,6 +50,8 @@ use jni::{
sys::jlongArray,
};
+/// ScanExec reads batches of data from Spark via JNI. The source of the scan could be a file
+/// scan or the result of reading a broadcast or shuffle exchange.
#[derive(Debug, Clone)]
pub struct ScanExec {
/// The ID of the execution context that owns this subquery. We use this ID to retrieve the JVM
@@ -64,7 +66,9 @@ pub struct ScanExec {
/// The input batch of input data. Used to determine the schema of the input data.
/// It is also used in unit test to mock the input data from JVM.
pub batch: Arc<Mutex<Option<InputBatch>>>,
+ /// Cache of expensive-to-compute plan properties
cache: PlanProperties,
+ /// Metrics collector
metrics: ExecutionPlanMetricsSet,
}
@@ -76,6 +80,11 @@ impl ScanExec {
data_types: Vec<DataType>,
) -> Result<Self, CometError> {
// Scan's schema is determined by the input batch, so we need to set it before execution.
+ // Note that we determine if arrays are dictionary-encoded based on the
+ // first batch. The array may be dictionary-encoded in some batches and not others, and
+ // ScanExec will cast arrays from all future batches to the type determined here, so we
+ // may end up either unpacking dictionary arrays or dictionary-encoding arrays.
+ // Dictionary-encoded primitive arrays are always unpacked.
let first_batch = if let Some(input_source) = input_source.as_ref() {
ScanExec::get_next(exec_context_id, input_source.as_obj())?
} else {
@@ -86,6 +95,8 @@ impl ScanExec {
let cache = PlanProperties::new(
EquivalenceProperties::new(schema),
+ // The partitioning is not important because we are not using DataFusion's
+ // query planner or optimizer
Partitioning::UnknownPartitioning(1),
ExecutionMode::Bounded,
);
@@ -109,11 +120,15 @@ impl ScanExec {
/// This is necessary since DataFusion doesn't handle dictionary array with values
/// being primitive type.
///
- /// TODO: revisit this once DF has imprved its dictionary type support. Ideally we shouldn't
+ /// TODO: revisit this once DF has improved its dictionary type support. Ideally we shouldn't
/// do this in Comet but rather let DF to handle it for us.
fn unpack_dictionary_type(dt: &DataType) -> DataType {
if let DataType::Dictionary(_, vt) = dt {
- if !matches!(vt.as_ref(), DataType::Utf8 | DataType::Binary) {
+ if !matches!(
+ vt.as_ref(),
+ DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary
+ ) {
+ // return the underlying data type
return vt.as_ref().clone();
}
}
@@ -128,13 +143,12 @@ impl ScanExec {
/// Pull next input batch from JVM.
pub fn get_next_batch(&mut self) -> Result<(), CometError> {
- let mut current_batch = self.batch.try_lock().unwrap();
-
if self.input_source.is_none() {
// This is a unit test. We don't need to call JNI.
return Ok(());
}
+ let mut current_batch = self.batch.try_lock().unwrap();
if current_batch.is_none() {
let next_batch = ScanExec::get_next(
self.exec_context_id,
@@ -153,8 +167,6 @@ impl ScanExec {
return Ok(InputBatch::EOF);
}
- let mut env = JVMClasses::get_env()?;
-
if iter.is_null() {
return Err(CometError::from(ExecutionError::GeneralError(format!(
"Null batch iterator object. Plan id: {}",
@@ -162,6 +174,7 @@ impl ScanExec {
))));
}
+ let mut env = JVMClasses::get_env()?;
let batch_object: JObject = unsafe {
jni_call!(&mut env,
comet_batch_iterator(iter).next() -> JObject)?
@@ -178,8 +191,6 @@ impl ScanExec {
let addresses = unsafe { env.get_array_elements(&batch_object, ReleaseMode::NoCopyBack)? };
- let mut inputs: Vec<ArrayRef> = vec![];
-
// First element is the number of rows.
let num_rows = unsafe { *addresses.as_ptr() as i64 };
@@ -197,9 +208,9 @@ impl ScanExec {
let num_arrays = array_num / 2;
let array_elements = unsafe { addresses.as_ptr().add(1) };
+ let mut inputs: Vec<ArrayRef> = Vec::with_capacity(num_arrays);
- let mut i: usize = 0;
- while i < num_arrays {
+ for i in 0..num_arrays {
let array_ptr = unsafe { *(array_elements.add(i * 2)) };
let schema_ptr = unsafe { *(array_elements.add(i * 2 + 1)) };
let array_data = ArrayData::from_spark((array_ptr, schema_ptr))?;
@@ -207,7 +218,6 @@ impl ScanExec {
// TODO: validate array input data
inputs.push(make_array(array_data));
- i += 1;
}
Ok(InputBatch::new(inputs, Some(num_rows as usize)))
@@ -310,22 +320,25 @@ impl DisplayAs for ScanExec {
}
/// A async-stream feeds input batch from `Scan` into DataFusion physical plan.
-struct ScanStream {
+struct ScanStream<'a> {
/// The `Scan` node producing input batches
scan: ScanExec,
/// Schema representing the data
schema: SchemaRef,
/// Metrics
baseline_metrics: BaselineMetrics,
+ /// Cast options
+ cast_options: CastOptions<'a>,
}
-impl ScanStream {
+impl<'a> ScanStream<'a> {
pub fn new(scan: ScanExec, schema: SchemaRef, partition: usize) -> Self {
let baseline_metrics = BaselineMetrics::new(&scan.metrics, partition);
Self {
scan,
schema,
baseline_metrics,
+ cast_options: CastOptions::default(),
}
}
@@ -337,16 +350,18 @@ impl ScanStream {
let schema_fields = self.schema.fields();
assert_eq!(columns.len(), schema_fields.len());
- // Cast if necessary
- let cast_options = CastOptions::default();
+ // Cast dictionary-encoded primitive arrays to regular arrays and cast
+ // Utf8/LargeUtf8/Binary arrays to dictionary-encoded if the schema is
+ // defined as dictionary-encoded and the data in this batch is not
+ // dictionary-encoded (could also be the other way around)
let new_columns: Vec<ArrayRef> = columns
.iter()
.zip(schema_fields.iter())
.map(|(column, f)| {
if column.data_type() != f.data_type() {
- cast_with_options(column, f.data_type(), &cast_options)
+ cast_with_options(column, f.data_type(), &self.cast_options)
} else {
- Ok(column.clone())
+ Ok(Arc::clone(column))
}
})
.collect::<Result<Vec<_>, _>>()?;
@@ -356,18 +371,18 @@ impl ScanStream {
}
}
-impl Stream for ScanStream {
+impl<'a> Stream for ScanStream<'a> {
type Item = DataFusionResult<RecordBatch>;
fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let mut timer = self.baseline_metrics.elapsed_compute().timer();
let mut scan_batch = self.scan.batch.try_lock().unwrap();
- timer.stop();
let input_batch = &*scan_batch;
let input_batch = if let Some(batch) = input_batch {
batch
} else {
+ timer.stop();
return Poll::Pending;
};
@@ -375,20 +390,20 @@ impl Stream for ScanStream {
InputBatch::EOF => Poll::Ready(None),
InputBatch::Batch(columns, num_rows) => {
self.baseline_metrics.record_output(*num_rows);
- let mut timer = self.baseline_metrics.elapsed_compute().timer();
let maybe_batch = self.build_record_batch(columns, *num_rows);
- timer.stop();
Poll::Ready(Some(maybe_batch))
}
};
*scan_batch = None;
+ timer.stop();
+
result
}
}
-impl RecordBatchStream for ScanStream {
+impl<'a> RecordBatchStream for ScanStream<'a> {
/// Get the schema
fn schema(&self) -> SchemaRef {
self.schema.clone()
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org