andygrove commented on code in PR #4791:
URL: https://github.com/apache/datafusion-comet/pull/4791#discussion_r4104874244
##########
native/shuffle/src/schema_align.rs:
##########
@@ -119,10 +141,16 @@ impl SchemaAlignExec {
expected_field.data_type()
);
}
- true
+ match (actual_field.data_type(), expected_field.data_type()) {
+ (DataType::LargeUtf8, DataType::Utf8) =>
ColumnAction::CastLargeStringToString,
Review Comment:
With the flag on, reusing `SchemaAlignExec` above the aggregate hits the
warning just above and logs `ShuffleWriter input schema mismatch on col[0]
'col_0': child produced LargeUtf8, catalyst declared Utf8. Inserting a cast;
please file the upstream function bug at .../issues/4515` on each executor. I
saw it in my local runs. There is no shuffle writer or upstream bug involved
here, and the module doc says this operator is enclosed by shuffle on purpose.
Would a small operator dedicated to the cast-back be cleaner? It could record
its own metrics, which would also help with the metrics problem sunchao raised
in `planner.rs`.
##########
native/shuffle/src/schema_align.rs:
##########
@@ -72,9 +75,27 @@ fn warn_dedup() -> &'static Mutex<HashSet<String>> {
pub struct SchemaAlignExec {
child: Arc<dyn ExecutionPlan>,
target_schema: SchemaRef,
+ column_actions: Arc<Vec<ColumnAction>>,
cache: Arc<PlanProperties>,
}
+#[derive(Debug, Clone)]
+enum ColumnAction {
Review Comment:
It looks like the rebase brought back the `ColumnAction` enum that #5138
removed from this file. #5138 routed `SchemaAlignStream` through
`cast_and_stamp_schema`, so every batch is checked against the target schema
and a failed cast names the operator and the column path. Here the choice is
made once at plan time. `Passthrough` columns are no longer checked per batch,
and the `Cast` arm returns a bare arrow error. This operator sits in front of
every native shuffle writer, so that applies with the flag off too. Could the
Large-to-small columns be handled first, with the result then handed to
`cast_and_stamp_schema`?
##########
spark/src/main/scala/org/apache/comet/CometConf.scala:
##########
@@ -455,6 +455,26 @@ object CometConf extends ShimCometConf {
.booleanConf
.createWithDefault(false)
+ val COMET_AGG_USE_LARGE_DATATYPES: ConfigEntry[Boolean] =
+ conf(s"$COMET_EXEC_CONFIG_PREFIX.aggregation.useLargeDataTypes")
Review Comment:
The other aggregate keys use `spark.comet.exec.aggregate.*`, and
`config_conventions.md` asks for feature flags to end in `.enabled`. This key
will be covered by the versioning policy once it ships, so could we settle on
something like `spark.comet.exec.aggregate.largeGroupKeys.enabled` now?
##########
spark/src/main/scala/org/apache/comet/CometConf.scala:
##########
@@ -455,6 +455,26 @@ object CometConf extends ShimCometConf {
.booleanConf
.createWithDefault(false)
+ val COMET_AGG_USE_LARGE_DATATYPES: ConfigEntry[Boolean] =
+ conf(s"$COMET_EXEC_CONFIG_PREFIX.aggregation.useLargeDataTypes")
+ .category(CATEGORY_EXEC)
+ .doc(
+ "When true, Comet wraps Utf8/Binary group-by expressions inside a Cast
to " +
+ "LargeUtf8/LargeBinary immediately before each native HashAggregate,
and wraps the " +
+ "aggregate's output in a Projection that casts those columns back to
Utf8/Binary. " +
+ "This promotes DataFusion's per-task group-key byte buffer from
`i32` offsets " +
+ "(2 GiB hard cap) to `i64` offsets, removing the `offset overflow,
buffer size > " +
+ "2147483647` failure that can hit CUBE / GROUPING SETS /
COUNT(DISTINCT) workloads " +
+ "where a single partition accumulates more than 2 GiB of distinct
string keys. " +
+ "Scan, shuffle, and the JVM FFI boundary remain Utf8/Binary, so
neither downstream " +
+ "Spark nor Comet's shuffle (which does not support LargeUtf8) is
affected. The cast " +
+ "is an offset-width promotion and only rebuilds the offset buffer
(the value bytes " +
+ "are shared), so per-batch overhead is O(rows) not O(bytes).
Defaults to false " +
+ "because the cap is only reachable for very large per-partition
group cardinalities; " +
+ "enable it when you see the offset-overflow error.")
+ .booleanConf
+ .createWithDefault(true)
Review Comment:
I asked in July whether defaulting this to `true` was intentional, and the
doc string still says it defaults to false. I measured it on a release build.
`SELECT count(*), sum(s) FROM (SELECT k, sum(v) s FROM t GROUP BY k)` over 8M
rows with 4M distinct 21-byte keys took 494 ms at the median with the flag on
and 450 ms with it off, about 10% slower over five alternating runs.
Low-cardinality and 200-byte keys were within noise. My laptop was busy, so
treat the numbers as rough. The overflow needs more than 2 GiB of distinct key
bytes in one task, and apache/datafusion#24704 tracks the real fix. Could this
default to `false`, with the tuning guide pointing at it for the `offset
overflow` error? Either way the doc needs updating. It describes the cast-back
as a `Projection` and says the overhead is O(rows), but as sunchao pointed out,
the cast-back copies every key byte.
##########
native/spark-expr/src/conversion_funcs/string.rs:
##########
@@ -172,7 +172,20 @@ impl TimeStampInfo {
}
pub(crate) fn is_df_cast_from_string_spark_compatible(to_type: &DataType) ->
bool {
- matches!(to_type, DataType::Binary)
+ // Utf8 -> Binary is a zero-copy reinterpret; Utf8 -> LargeUtf8 is a pure
+ // offset-width widening handled by arrow's cast_byte_container. Both are
+ // Spark-equivalent (Spark's StringType and BinaryType are indifferent to
+ // Arrow's offset width).
+ matches!(to_type, DataType::Binary | DataType::LargeUtf8)
+}
+
+pub(crate) fn is_df_cast_from_large_string_spark_compatible(to_type:
&DataType) -> bool {
Review Comment:
I can't find a way to reach these new arms, or the matching ones in
`cast.rs`. `types.proto` has no Large type ids, so the serde never emits a
`Cast` to or from one. The Parquet schema adapter's casts already go through
the `is_adapting_schema` branch in `cast_array`. The comment also says
`SchemaAlignExec` pre-splits arrays before they get here, but `SchemaAlignExec`
builds its arrays directly and never calls this cast. Could these arms be
removed from this PR?
##########
spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala:
##########
@@ -578,10 +578,15 @@ object Utils extends CometTypeShim with Logging {
}
def getFieldVector(valueVector: ValueVector, reason: String): FieldVector = {
- if (isSupportedFieldVector(valueVector)) {
- valueVector.asInstanceOf[FieldVector]
- } else {
- throw new SparkException(s"Unsupported Arrow Vector for $reason:
${valueVector.getClass}")
+ valueVector match {
+ case v if isSupportedFieldVector(v) =>
+ v.asInstanceOf[FieldVector]
+ // Accepted here but left out of isSupportedFieldVector, which
isArrowBacked uses to keep
Review Comment:
Which path needs `getFieldVector` to accept `LargeVarCharVector` now that
the aggregate casts back natively? As far as I can tell, the only Large vectors
on the JVM side come from PyArrow UDFs returning `large_string`. So this
changes native C2R export and broadcast serialization for that path without a
test. It also makes the `isSupportedFieldVector` doc, the `isArrowBacked`
comment and the `UtilsSuite` test comment wrong, since all three assume
`getFieldVector` rejects these vectors. Could this be dropped here, or moved to
its own PR with a PyArrow test? The same question applies to the `LargeUtf8`
passthrough at `columnar_to_row.rs:1041`.
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -3344,4 +3350,111 @@ class CometAggregateSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
+ // The test to reproduce `offset overflow` for aggregation queries, when
interim data
+ // get exploded 100x comparing to initial input size.
+ // It is not supposed to run on CI as the test requires significant RAM to
succeed
+ ignore("CUBE(9) + COUNT(DISTINCT) wide Utf8 keys: useLargeDataTypes
preserves correctness") {
Review Comment:
This test is `ignore`d, so CI never runs it. The `sparkConf` comment at line
60 still describes the off-heap and memory pool settings that 7252adb853
removed. This test now sets the pool through `withSQLConf`, which that comment
says won't take effect. You mentioned in July that CI passes with the flag on.
That covers the common path, but nothing exercises the split in
`compute_row_ranges` or checks the metrics. Could the byte cap be a parameter
so a Rust unit test can split a small batch? And could this test and the stale
comment be replaced with small tests that run? If the default moves to `false`,
those tests would also need to turn the flag on explicitly.
##########
native/core/src/execution/operators/shuffle_scan.rs:
##########
@@ -210,15 +210,20 @@ impl ShuffleScanExec {
let num_rows = batch.num_rows();
- // Extract column arrays, unpacking any dictionary-encoded columns.
- // Native shuffle may dictionary-encode string/binary columns for
efficiency,
- // but downstream DataFusion operators expect the value types
declared in the
- // schema (e.g. Utf8, not Dictionary<Int32, Utf8>).
+ // Coerce each decoded column to the catalyst-declared type:
Review Comment:
Since the aggregate now casts back before anything above it sees the batch,
I don't think this feature can put `LargeUtf8` into a shuffle block. If one
did, the remote read path would reject it, because `remote_schema.rs` treats
`LargeUtf8` against `Utf8` as an incompatible type. On the local path,
`ShuffleScanStream::poll_next` already reconciles every column through
`cast_and_stamp_schema`, so the extra cast here duplicates that. Could this go
back to unpacking dictionaries only? Returning an error instead of the old
`expect` is a good change and worth keeping.
##########
native/core/src/execution/planner.rs:
##########
@@ -1451,6 +1498,36 @@ impl PhysicalPlanner {
)?,
);
+ // Cast promoted group columns back to their original
Utf8/Binary type
+ // so LargeUtf8/LargeBinary never crosses the FFI boundary,
JVM columnar
+ // shuffle, or the Spark consumer path (all of which reject
Large*).
+ // Uses `SchemaAlignExec` (not a plain `ProjectionExec` +
`CastExpr`)
+ // because arrow's cast kernel rejects any single Large* array
whose
+ // value bytes exceed `i32::MAX` -- which is precisely the
regime this
+ // flag is used in. `SchemaAlignExec` splits each batch by row
ranges
+ // so every emitted small-offset chunk fits under 2 GiB.
+ let aggregate = if use_large &&
group_reverts.iter().any(Option::is_some) {
+ let agg_schema = aggregate.schema();
+ let target_fields: Vec<Field> = agg_schema
+ .fields()
+ .iter()
+ .enumerate()
+ .map(|(idx, f)| {
+ let target_dt = group_reverts
+ .get(idx)
+ .and_then(|r| r.clone())
+ .unwrap_or_else(|| f.data_type().clone());
+ Field::new(f.name(), target_dt, f.is_nullable())
+ .with_metadata(f.metadata().clone())
+ })
+ .collect();
+ let target_schema: SchemaRef =
Arc::new(Schema::new(target_fields));
+ SchemaAlignExec::try_new_or_passthrough(aggregate,
&target_schema)
Review Comment:
I can confirm this end to end with a full build. `SELECT k, sum(v) FROM t
GROUP BY k` on a string key shows the partial `CometHashAggregateExec` with
`output_rows=0` and `elapsed_compute=0` when the flag is on, against 20000 rows
and about 50 ms with it off. The final aggregate's `elapsed_compute` drops from
about 10 ms to 17 µs. When fixing it, note that `to_native_metric_node` skips
`output_rows` from `additional_native_plans`, so registering the
`AggregateExec` through `SparkPlan::new_with_additional` also needs the wrapper
to record its own `output_rows`. A test asserting these metrics on a
string-keyed aggregate would catch this. The one in `CometExecSuite` uses a
numeric key.
--
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]
