comphead commented on code in PR #6107:
URL: https://github.com/apache/datafusion-comet/pull/6107#discussion_r4414766322
##########
native/operators/src/explode.rs:
##########
@@ -346,6 +356,82 @@ impl PendingInput {
.as_ref()
.map(|lens| lens.slice(self.row_offset, rows))
}
+
+ fn next_chunk(
+ &mut self,
+ batch_size: usize,
+ list_columns: &[ListUnnest],
+ ) -> Result<(RecordBatch, Option<Int64Array>)> {
+ if let Some(limit) = self.split_limit {
+ let row_length =
self.output_lens.as_ref().unwrap().value(self.row_offset) as usize;
+ if row_length > limit {
+ let length = limit.min(row_length - self.element_offset);
+ let row = self.batch.slice(self.row_offset, 1);
+ let mut columns = row.columns().to_vec();
+ for unnest in list_columns {
+ let list =
columns[unnest.index_in_input_schema].as_list::<i32>();
+ let DataType::List(field) = list.data_type() else {
+ unreachable!();
+ };
+ let available = if list.is_null(0) {
+ 0
+ } else {
+ list.value_length(0) as usize
+ };
+ let skip = self.element_offset.min(available);
+ let start = list.value_offsets()[0] + skip as i32;
+ let end = start + length.min(available - skip) as i32;
+ // Narrow the entries and parallel positions together,
preserving their
+ // original values and sharing their buffers. Shorter
lists still pad.
+ columns[unnest.index_in_input_schema] =
Arc::new(ListArray::new(
+ Arc::clone(field),
+ OffsetBuffer::new(vec![start, end].into()),
+ Arc::clone(list.values()),
+ list.nulls().cloned(),
+ ));
+ }
+ self.element_offset += length;
+ if self.element_offset == row_length {
+ self.row_offset += 1;
+ self.element_offset = 0;
+ }
+ return Ok((
+ RecordBatch::try_new(row.schema(), columns)?,
+ Some(Int64Array::from(vec![length as i64])),
+ ));
+ }
+ }
+ let rows = self.next_chunk_rows(batch_size);
+ let chunk = self.batch.slice(self.row_offset, rows);
+ let lengths = self.chunk_lengths(rows);
+ self.row_offset += rows;
+ Ok((chunk, lengths))
+ }
+}
+
+/// If every row repeats at most R times, no copied buffer or child offset
span exceeds R
+/// times its entire source span. Count elements too: List<Null> has no child
value buffer.
+fn repetition_limit(batch: &RecordBatch, list_columns: &[ListUnnest]) -> usize
{
+ fn max_span(data: &ArrayData) -> usize {
+ std::iter::once(data.len())
+ .chain(data.buffers().iter().map(|buffer| buffer.len()))
+ .chain(data.child_data().iter().map(max_span))
+ .max()
+ .unwrap_or(0)
+ }
+ let span = batch
+ .columns()
+ .iter()
+ .enumerate()
+ .filter(|(index, _)| {
+ !list_columns
+ .iter()
+ .any(|list| list.index_in_input_schema == *index)
+ })
+ .map(|(_, column)| max_span(&column.to_data()))
+ .max()
+ .unwrap_or(0);
+ (i32::MAX as usize / span.max(1)).max(1)
Review Comment:
`span` is the largest buffer of any carried column across the whole input
batch, but what can overflow is the bytes of the rows in one chunk times their
repetition. Take an 8,192-row batch carrying a 1 KiB string. Its values buffer
is 8 MiB, so the limit is 255, and every row with more than 255 elements is
then built alone in 255-element pieces, with `next_chunk_rows` also ending the
chunk before such a row. On main, 300-element arrays with that payload are
built about 27 rows at a time, roughly 8 MB of copied strings, far from the
offset limit. `split_limit` is set for every depth-1 `List`, so this changes
array `explode` too, and fixed-width carried columns count toward `span`
although they cannot overflow an offset. Could the limit come from the
offset-addressed bytes of the rows in the chunk, or apply only when a row's own
carried bytes times its length can overflow? I derived this from the code and
have not measured it. A bench case with a 1 KiB carried string and a few hundr
ed elements per row would show it.
##########
native/operators/src/explode.rs:
##########
@@ -1276,36 +1374,164 @@ fn create_take_indices(
/// c2: 'a', 'b', 'c', 'c', 'c', null, 'd', 'd'
/// ```
///
-/// The `repeat_mask` determines whether an array's values are repeated or
replaced with nulls.
-/// For example, if the `repeat_mask` is:
-///
-/// ```ignore
-/// [true, false]
-/// ```
-///
-/// The final batch will look like:
-///
-/// ```ignore
-/// c1: 1, null, 2, 3, 4, null, 5, 6 // Repeated using `indices`
-/// c2: null, null, null, null, null, null, null, null // Replaced with nulls
-fn repeat_arrs_from_indices(
+fn repeat_arrs(
batch: &[ArrayRef],
- indices: &PrimitiveArray<Int64Type>,
+ lengths: &PrimitiveArray<Int64Type>,
+ total_length: usize,
repeat_mask: &[bool],
) -> Result<Vec<Arc<dyn Array>>> {
+ let mut indices = None;
batch
.iter()
.zip(repeat_mask.iter())
.map(|(arr, &repeat)| {
if repeat {
+ if let Some(strings) = arr.as_string_opt::<i32>() {
+ if let Some(strings) = repeat_short_strings(strings,
lengths, total_length) {
+ return Ok(Arc::new(strings) as ArrayRef);
+ }
+ }
+ if arr.null_count() == 0 && total_length >=
arr.len().saturating_mul(5) {
+ let repeated = arrow::array::downcast_primitive_array! {
+ arr => { Some(Arc::new(repeat_primitive(arr, lengths,
total_length)) as ArrayRef) }
+ _ => None,
+ };
+ if let Some(repeated) = repeated {
+ return Ok(repeated);
+ }
+ }
+ let indices =
+ indices.get_or_insert_with(|| create_take_indices(lengths,
total_length));
Ok(kernels::take::take(arr, indices, None)?)
} else {
- Ok(new_null_array(arr.data_type(), arr.len()))
+ Ok(Arc::clone(arr))
}
})
.collect()
}
+/// Repeat non-null primitives by input runs without allocating gather indices.
+fn repeat_primitive<T: arrow::datatypes::ArrowPrimitiveType>(
+ array: &PrimitiveArray<T>,
+ lengths: &Int64Array,
+ total_length: usize,
+) -> PrimitiveArray<T> {
+ let mut values = Vec::with_capacity(total_length);
+ for (&value, &repeat) in array.values().iter().zip(lengths.values()) {
+ values.resize(values.len() + repeat as usize, value);
+ }
+ PrimitiveArray::new(values.into(),
None).with_data_type(array.data_type().clone())
+}
+
+/// Repeat short Utf8 values using the input runs rather than gathering every
output value.
+/// Batches dominated by long strings or with less than twofold expansion use
Arrow's take instead.
+fn repeat_short_strings(
Review Comment:
`repeat_short_strings`, with its `StringArray::new_unchecked`, and
`repeat_primitive` change how every `explode` repeats carried columns, arrays
included, and map support does not need them. There are no current-head numbers
for them, and the new `explode_chained_carry` and `fan_out_*` Criterion groups
have no posted results. Could they move to their own PR with before and after
numbers from `benches/explode.rs`? That keeps this PR to map support and lets
the `unsafe` path be reverted on its own, as andygrove asked for the
`ParquetReadSuite` change.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1444,6 +1454,29 @@ case class CometExecRule(session: SparkSession,
queryStagePrep: Boolean = false)
visit(plan, if (producesObjects(plan)) Some("code reading Dataset.rdd")
else None)
}
+ // Spark fuses a generator chain without materializing its carried columns
between generators.
+ // Keep that path until native map expansion can avoid the extra repetition
too. Tags survive
+ // the child copies made by bottom-up conversion; exchanges end the shared
pipeline.
+ private def tagChainedMapGenerators(plan: SparkPlan): Unit = {
+ val enabled = CometConf.COMET_EXEC_EXPLODE_ENABLED.get(conf)
+ def visit(op: SparkPlan, generatorAbove: Boolean): Unit = op match {
+ case generate: GenerateExec =>
+ generate.unsetTagValue(CometExecRule.CHAINED_MAP_GENERATOR)
+ generate.generator match {
+ case explode: ExplodeBase
+ if enabled && generatorAbove &&
explode.child.dataType.isInstanceOf[MapType] =>
Review Comment:
This keeps a map generator on Spark whenever another generator sits above
it, while an array generator in the same position stays native, and `array-only
generator chain stays native` asserts that. The stated reason is carried-column
repetition between generators, which costs the same whether the lower generator
reads a map or an array. The description calls the 97 ms against 90 ms
comparison historical, and `c4525cb85f` has since changed how carried values
repeat. Could you post nested numbers at this head for the map chain and the
equivalent `array<struct>` chain? If the gap is gone, the rule, its tests and
the `operators.md` paragraph can go. If it remains, it is a chained-generator
issue rather than a map one. The benchmark also replaced the native
`array<struct>` depth-8 chain with this map chain, which now times Spark
generators, so keeping both cases would answer this.
##########
spark/src/test/resources/sql-tests/expressions/array/explode.sql:
##########
@@ -339,7 +339,7 @@ INSERT INTO test_explode_map VALUES
(2, map()),
(3, NULL)
-query expect_fallback(Comet only supports explode/explode_outer for arrays,
not maps)
+query
Review Comment:
This and the `posexplode(m)` case in `posexplode.sql` were fallback checks,
and as plain queries they now repeat `explode_maps` in `map/explode.sql`, which
has the same entries, empty and NULL shapes for all four generators. Could they
be dropped? `with batches containing only null and empty maps` in
`CometGenerateExecSuite` also looks like a fit for that fixture with `--
Config: spark.comet.batchSize=4`. I have not run that conversion.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
