From ef149172c9b5ab3f61734935cf961cc9ba408fab Mon Sep 17 00:00:00 2001 From: xylaaaaa <2392805527@qq.com> Date: Tue, 14 Jul 2026 00:23:30 +0800 Subject: [PATCH] [fix](be) Avoid materializing complex columns for external count ### What problem does this PR solve? Issue Number: None Related PR: None Problem Summary: COUNT(nullable complex column) over Parquet or ORC file TVFs cannot use COUNT(*) row metadata because top-level NULL values must be excluded. With FileScannerV2, the query therefore decoded the full complex column through Arrow and then materialized Doris columns, retaining decoded binary chunks and causing excessive peak memory or intermittent OOM. This change introduces an explicit COUNT_NON_NULL operation and target slot from FE to BE. Parquet counts top-level shape from bounded definition/repetition-level batches without Arrow builders or Doris value columns, while ORC uses non-null stripe statistics. It also emits valid non-NULL synthetic target rows and releases file-block ownership before mutable output conversion to avoid unnecessary COW copies. ### Release note Reduce memory usage for COUNT on nullable columns read from Parquet and ORC file TVFs with FileScannerV2. ### Check List (For Author) - Test: Unit Test - FE PhysicalStorageLayerAggregateTest: 4 tests passed - BE targeted Parquet, ORC, and TableReader tests: 14 tests passed - HDFS TVF regression assertion added but not run locally because this worktree has no Doris cluster or external HDFS test environment - Behavior changed: Yes. Eligible nullable file-column counts use explicit COUNT_NON_NULL pushdown instead of materializing the full column. - Does this need documentation: No --- be/src/exec/operator/scan_operator.cpp | 8 + be/src/exec/operator/scan_operator.h | 5 + be/src/exec/scan/file_scanner_v2.cpp | 14 +- be/src/format_v2/orc/orc_reader.cpp | 8 +- be/src/format_v2/parquet/parquet_reader.cpp | 17 +- .../parquet/reader/column_reader.cpp | 197 +++++++++++++++++- .../format_v2/parquet/reader/column_reader.h | 1 - .../parquet/reader/parquet_leaf_reader.h | 11 +- be/src/format_v2/table_reader.cpp | 4 + be/src/format_v2/table_reader.h | 112 +++++++--- be/test/format_v2/orc/orc_reader_test.cpp | 2 +- .../format_v2/parquet/parquet_reader_test.cpp | 14 +- be/test/format_v2/table_reader_test.cpp | 97 ++++++++- .../apache/doris/datasource/FileScanNode.java | 3 + .../translator/PhysicalPlanTranslator.java | 37 +++- .../translator/PlanTranslatorContext.java | 9 + .../apache/doris/nereids/rules/RuleType.java | 2 + .../implementation/AggregateStrategies.java | 96 ++++++++- .../PhysicalStorageLayerAggregate.java | 96 +++++++-- .../org/apache/doris/planner/PlanNode.java | 5 + .../PhysicalStorageLayerAggregateTest.java | 66 +++++- gensrc/thrift/PlanNodes.thrift | 5 +- .../tvf/test_hdfs_parquet_group0.groovy | 9 + 23 files changed, 726 insertions(+), 92 deletions(-) diff --git a/be/src/exec/operator/scan_operator.cpp b/be/src/exec/operator/scan_operator.cpp index a9e92d08e08613..18f8d9ce453731 100644 --- a/be/src/exec/operator/scan_operator.cpp +++ b/be/src/exec/operator/scan_operator.cpp @@ -1046,6 +1046,11 @@ TPushAggOp::type ScanLocalState::get_push_down_agg_type() { return _parent->cast()._push_down_agg_type; } +template +std::optional ScanLocalState::get_count_non_null_slot_id() { + return _parent->cast()._count_non_null_slot_id; +} + template int64_t ScanLocalState::limit_per_scanner() { return _parent->cast()._limit_per_scanner; @@ -1231,6 +1236,9 @@ Status ScanOperatorX::init(const TPlanNode& tnode, RuntimeState* } else { _push_down_agg_type = TPushAggOp::type::NONE; } + if (tnode.__isset.file_scan_node && tnode.file_scan_node.__isset.count_non_null_slot_id) { + _count_non_null_slot_id = tnode.file_scan_node.count_non_null_slot_id; + } if (tnode.__isset.topn_filter_source_node_ids) { _topn_filter_source_node_ids = tnode.topn_filter_source_node_ids; diff --git a/be/src/exec/operator/scan_operator.h b/be/src/exec/operator/scan_operator.h index 8b12ccf0bc1195..422767d1a71acc 100644 --- a/be/src/exec/operator/scan_operator.h +++ b/be/src/exec/operator/scan_operator.h @@ -18,6 +18,7 @@ #pragma once #include +#include #include #include @@ -74,6 +75,7 @@ class ScanLocalStateBase : public PipelineXLocalState<> { virtual void set_scan_ranges(RuntimeState* state, const std::vector& scan_ranges) = 0; virtual TPushAggOp::type get_push_down_agg_type() = 0; + virtual std::optional get_count_non_null_slot_id() = 0; // If scan operator is serial operator(like topn), its real parallelism is 1. // Otherwise, its real parallelism is query_parallel_instance_num. @@ -252,6 +254,7 @@ class ScanLocalState : public ScanLocalStateBase { const std::vector& scan_ranges) override {} TPushAggOp::type get_push_down_agg_type() override; + std::optional get_count_non_null_slot_id() override; std::vector execution_dependencies() override { if (_filter_dependencies.empty()) { @@ -386,6 +389,7 @@ class ScanOperatorX : public OperatorX { } TPushAggOp::type get_push_down_agg_type() { return _push_down_agg_type; } + std::optional get_count_non_null_slot_id() const { return _count_non_null_slot_id; } DataDistribution required_data_distribution(RuntimeState* /*state*/) const override { if (OperatorX::is_serial_operator()) { @@ -451,6 +455,7 @@ class ScanOperatorX : public OperatorX { std::vector _runtime_filter_descs; TPushAggOp::type _push_down_agg_type; + std::optional _count_non_null_slot_id; // Record the value of the aggregate function 'count' from doris's be int64_t _push_down_count = -1; diff --git a/be/src/exec/scan/file_scanner_v2.cpp b/be/src/exec/scan/file_scanner_v2.cpp index 6700ed71276f45..9fc201770cc5eb 100644 --- a/be/src/exec/scan/file_scanner_v2.cpp +++ b/be/src/exec/scan/file_scanner_v2.cpp @@ -448,6 +448,16 @@ Status FileScannerV2::_init_table_reader(const TFileRangeDesc& range) { VExprContextSPtrs table_conjuncts; RETURN_IF_ERROR(_build_table_conjuncts(&table_conjuncts)); + std::optional count_non_null_global_index; + if (_local_state->get_count_non_null_slot_id().has_value()) { + const auto target = + _slot_id_to_global_index.find(*_local_state->get_count_non_null_slot_id()); + if (target == _slot_id_to_global_index.end()) { + return Status::InternalError("Unknown COUNT_NON_NULL target slot {}", + *_local_state->get_count_non_null_slot_id()); + } + count_non_null_global_index = target->second; + } RETURN_IF_ERROR(_table_reader->init({ .projected_columns = _projected_columns, .conjuncts = std::move(table_conjuncts), @@ -458,6 +468,7 @@ Status FileScannerV2::_init_table_reader(const TFileRangeDesc& range) { .scanner_profile = _local_state->scanner_profile(), .file_slot_descs = &_file_slot_descs, .push_down_agg_type = _local_state->get_push_down_agg_type(), + .count_non_null_global_index = count_non_null_global_index, .condition_cache_digest = _local_state->get_condition_cache_digest(), })); return Status::OK(); @@ -800,7 +811,8 @@ bool FileScannerV2::_should_run_adaptive_batch_size() const { // COUNT pushdown emits synthetic rows from file metadata and does not materialize file columns, // so there is no useful row-width sample to learn from. return _block_size_predictor != nullptr && - _local_state->get_push_down_agg_type() != TPushAggOp::type::COUNT; + _local_state->get_push_down_agg_type() != TPushAggOp::type::COUNT && + _local_state->get_push_down_agg_type() != TPushAggOp::type::COUNT_NON_NULL; } size_t FileScannerV2::_predict_reader_batch_rows() { diff --git a/be/src/format_v2/orc/orc_reader.cpp b/be/src/format_v2/orc/orc_reader.cpp index b728c6fbedb700..87404bab63c83f 100644 --- a/be/src/format_v2/orc/orc_reader.cpp +++ b/be/src/format_v2/orc/orc_reader.cpp @@ -1951,6 +1951,7 @@ Status OrcReader::get_aggregate_result(const format::FileAggregateRequest& reque result->count = 0; result->columns.clear(); if (request.agg_type != TPushAggOp::type::COUNT && + request.agg_type != TPushAggOp::type::COUNT_NON_NULL && request.agg_type != TPushAggOp::type::MINMAX) { return Status::NotSupported("Unsupported ORC aggregate pushdown type {}", request.agg_type); } @@ -1984,8 +1985,13 @@ Status OrcReader::get_aggregate_result(const format::FileAggregateRequest& reque result->count += cast_set(stripe_information->getNumberOfRows()); } - if (request.agg_type == TPushAggOp::type::COUNT) { + if (request.agg_type == TPushAggOp::type::COUNT || + request.agg_type == TPushAggOp::type::COUNT_NON_NULL) { if (request.columns.empty()) { + if (request.agg_type == TPushAggOp::type::COUNT_NON_NULL) { + return Status::InvalidArgument( + "ORC COUNT_NON_NULL pushdown requires one count column"); + } return Status::OK(); } if (request.columns.size() != 1) { diff --git a/be/src/format_v2/parquet/parquet_reader.cpp b/be/src/format_v2/parquet/parquet_reader.cpp index 753b3628bfa19b..c06ebf0d157291 100644 --- a/be/src/format_v2/parquet/parquet_reader.cpp +++ b/be/src/format_v2/parquet/parquet_reader.cpp @@ -593,6 +593,7 @@ Status ParquetReader::get_aggregate_result(const format::FileAggregateRequest& r result->count = 0; result->columns.clear(); if (request.agg_type != TPushAggOp::type::COUNT && + request.agg_type != TPushAggOp::type::COUNT_NON_NULL && request.agg_type != TPushAggOp::type::MINMAX) { return Status::NotSupported("Unsupported parquet aggregate pushdown type {}", request.agg_type); @@ -605,8 +606,13 @@ Status ParquetReader::get_aggregate_result(const format::FileAggregateRequest& r DORIS_CHECK(row_group_metadata != nullptr); result->count += row_group_metadata->num_rows(); } - if (request.agg_type == TPushAggOp::type::COUNT) { + if (request.agg_type == TPushAggOp::type::COUNT || + request.agg_type == TPushAggOp::type::COUNT_NON_NULL) { if (request.columns.empty()) { + if (request.agg_type == TPushAggOp::type::COUNT_NON_NULL) { + return Status::InvalidArgument( + "Parquet COUNT_NON_NULL pushdown requires one count column"); + } return Status::OK(); } if (request.columns.size() != 1) { @@ -654,11 +660,10 @@ Status ParquetReader::get_aggregate_result(const format::FileAggregateRequest& r while (range_rows_read < selected_range.length) { const int64_t batch_rows = std::min(_batch_size, selected_range.length - range_rows_read); - // COUNT(col) only needs the top-level NULL state. The shape reader loads - // def/rep levels from one representative leaf and does not build value_indices - // or values_column. MAP chooses the key leaf; ARRAY/STRUCT may choose a string - // leaf, but the levels-only protocol still avoids Doris-side string - // materialization for that leaf. + // COUNT(col) only needs the top-level NULL state. The shape reader keeps one + // def/rep pair per top-level row and discards physical values. Binary leaves + // use lightweight ByteArray cursors, so neither Arrow binary builders nor + // Doris string columns are materialized. RETURN_IF_ERROR(_stop_status_if_requested( shape_reader->load_nested_levels_batch(batch_rows))); _record_scan_rows(batch_rows); diff --git a/be/src/format_v2/parquet/reader/column_reader.cpp b/be/src/format_v2/parquet/reader/column_reader.cpp index 352fbbd7c3d215..5072459be217ac 100644 --- a/be/src/format_v2/parquet/reader/column_reader.cpp +++ b/be/src/format_v2/parquet/reader/column_reader.cpp @@ -22,6 +22,7 @@ #include #include +#include #include #include #include @@ -30,6 +31,7 @@ #include #include +#include "common/cast_set.h" #include "core/data_type/data_type_array.h" #include "core/data_type/data_type_map.h" #include "core/data_type/data_type_nullable.h" @@ -133,6 +135,173 @@ bool supports_nested_scalar_record_reader(const ParquetColumnSchema& column_sche return true; } +class CountShapeColumnReader final : public ParquetColumnReader { +public: + CountShapeColumnReader(const ParquetColumnSchema& column_schema, + std::shared_ptr<::parquet::ColumnReader> column_reader, + ParquetColumnReaderProfile profile) + : ParquetColumnReader(column_schema, column_schema.type, profile), + _column_reader(std::move(column_reader)), + _descriptor(column_schema.descriptor) { + DORIS_CHECK(_column_reader != nullptr); + DORIS_CHECK(_descriptor != nullptr); + } + + Status read(int64_t rows, MutableColumnPtr&, int64_t*) override { + return Status::NotSupported("COUNT shape reader cannot materialize {} rows for column {}", + rows, _name); + } + + Status skip(int64_t rows) override { + std::vector def_levels; + std::vector rep_levels; + RETURN_IF_ERROR(_consume_top_level_rows(rows, &def_levels, &rep_levels)); + update_reader_skip_rows(rows); + return Status::OK(); + } + + Status load_nested_levels_batch(int64_t rows) override { + RETURN_IF_ERROR(_consume_top_level_rows(rows, &_definition_levels, &_repetition_levels)); + update_reader_read_rows(rows); + return Status::OK(); + } + + const std::vector& nested_definition_levels() const override { + return _definition_levels; + } + + const std::vector& nested_repetition_levels() const override { + return _repetition_levels; + } + + int64_t nested_levels_written() const override { + return cast_set(_definition_levels.size()); + } + + bool is_or_has_repeated_child() const override { + return _descriptor->max_repetition_level() > 0; + } + +private: + struct TopLevelShape { + int16_t definition_level = 0; + int16_t repetition_level = 0; + }; + + template + int64_t _read_typed_batch(int64_t batch_size, int16_t* def_levels, int16_t* rep_levels, + int64_t* values_read) { + auto* typed_reader = + dynamic_cast<::parquet::TypedColumnReader*>(_column_reader.get()); + DORIS_CHECK(typed_reader != nullptr); + using ValueType = typename PhysicalType::c_type; + auto values = std::make_unique(cast_set(batch_size)); + return typed_reader->ReadBatch(batch_size, def_levels, rep_levels, values.get(), + values_read); + } + + int64_t _read_physical_batch(int64_t batch_size, int16_t* def_levels, int16_t* rep_levels, + int64_t* values_read) { + switch (_column_reader->type()) { + case ::parquet::Type::BOOLEAN: + return _read_typed_batch<::parquet::BooleanType>(batch_size, def_levels, rep_levels, + values_read); + case ::parquet::Type::INT32: + return _read_typed_batch<::parquet::Int32Type>(batch_size, def_levels, rep_levels, + values_read); + case ::parquet::Type::INT64: + return _read_typed_batch<::parquet::Int64Type>(batch_size, def_levels, rep_levels, + values_read); + case ::parquet::Type::INT96: + return _read_typed_batch<::parquet::Int96Type>(batch_size, def_levels, rep_levels, + values_read); + case ::parquet::Type::FLOAT: + return _read_typed_batch<::parquet::FloatType>(batch_size, def_levels, rep_levels, + values_read); + case ::parquet::Type::DOUBLE: + return _read_typed_batch<::parquet::DoubleType>(batch_size, def_levels, rep_levels, + values_read); + case ::parquet::Type::BYTE_ARRAY: + return _read_typed_batch<::parquet::ByteArrayType>(batch_size, def_levels, rep_levels, + values_read); + case ::parquet::Type::FIXED_LEN_BYTE_ARRAY: + return _read_typed_batch<::parquet::FLBAType>(batch_size, def_levels, rep_levels, + values_read); + case ::parquet::Type::UNDEFINED: + break; + } + throw ::parquet::ParquetException("Unsupported physical type for COUNT shape reader"); + } + + Status _load_more_shapes() { + constexpr int64_t LEVEL_BATCH_SIZE = 8192; + std::vector def_levels(LEVEL_BATCH_SIZE); + std::vector rep_levels(LEVEL_BATCH_SIZE); + int64_t levels_read = 0; + try { + int64_t values_read = 0; + levels_read = _read_physical_batch(LEVEL_BATCH_SIZE, def_levels.data(), + rep_levels.data(), &values_read); + } catch (const ::parquet::ParquetException& e) { + return Status::Corruption("Failed to read parquet shape for column {}: {}", _name, + e.what()); + } catch (const std::exception& e) { + return Status::InternalError("Failed to read parquet shape for column {}: {}", _name, + e.what()); + } + if (levels_read <= 0 || levels_read > LEVEL_BATCH_SIZE) { + return Status::Corruption("Invalid parquet shape batch for column {}: levels={}", _name, + levels_read); + } + + const bool repeated = _descriptor->max_repetition_level() > 0; + for (int64_t level_idx = 0; level_idx < levels_read; ++level_idx) { + const int16_t repetition_level = repeated ? rep_levels[level_idx] : 0; + if (repetition_level != 0) { + continue; + } + const int16_t definition_level = + _descriptor->max_definition_level() > 0 ? def_levels[level_idx] : 0; + _pending_shapes.push_back({definition_level, repetition_level}); + } + return Status::OK(); + } + + Status _consume_top_level_rows(int64_t rows, std::vector* def_levels, + std::vector* rep_levels) { + DORIS_CHECK(rows >= 0); + DORIS_CHECK(def_levels != nullptr); + DORIS_CHECK(rep_levels != nullptr); + def_levels->clear(); + rep_levels->clear(); + def_levels->reserve(cast_set(rows)); + rep_levels->reserve(cast_set(rows)); + while (def_levels->size() < cast_set(rows)) { + if (_pending_shapes.empty()) { + if (!_column_reader->HasNext()) { + return Status::Corruption( + "Parquet COUNT shape reader reached EOF after {} of {} rows for column " + "{}", + def_levels->size(), rows, _name); + } + RETURN_IF_ERROR(_load_more_shapes()); + continue; + } + const auto shape = _pending_shapes.front(); + _pending_shapes.pop_front(); + def_levels->push_back(shape.definition_level); + rep_levels->push_back(shape.repetition_level); + } + return Status::OK(); + } + + std::shared_ptr<::parquet::ColumnReader> _column_reader; + const ::parquet::ColumnDescriptor* _descriptor = nullptr; + std::deque _pending_shapes; + std::vector _definition_levels; + std::vector _repetition_levels; +}; + } // namespace Status ParquetColumnReader::skip(int64_t rows) { @@ -486,12 +655,12 @@ Status ParquetColumnReaderFactory::create(const ParquetColumnSchema& column_sche Status ParquetColumnReaderFactory::create_count_shape_reader( const ParquetColumnSchema& column_schema, const format::LocalColumnIndex* projection, std::unique_ptr* reader) const { - return create_count_shape_reader_impl(column_schema, projection, false, reader); + return create_count_shape_reader_impl(column_schema, projection, reader); } Status ParquetColumnReaderFactory::create_count_shape_reader_impl( const ParquetColumnSchema& column_schema, const format::LocalColumnIndex* projection, - bool is_nested, std::unique_ptr* reader) const { + std::unique_ptr* reader) const { if (reader == nullptr) { return Status::InvalidArgument("reader is null"); } @@ -501,7 +670,23 @@ Status ParquetColumnReaderFactory::create_count_shape_reader_impl( return Status::InvalidArgument("Parquet COUNT projection is invalid for column {}", column_schema.name); } - return create_scalar_column_reader(column_schema, is_nested, false, reader); + DORIS_CHECK(_row_group != nullptr); + DORIS_CHECK(column_schema.leaf_column_id >= 0); + try { + auto physical_reader = _row_group->Column(column_schema.leaf_column_id); + DORIS_CHECK(physical_reader != nullptr); + *reader = std::make_unique( + column_schema, std::move(physical_reader), _column_reader_profile); + return Status::OK(); + } catch (const ::parquet::ParquetException& e) { + return Status::Corruption( + "Failed to create parquet COUNT shape reader for column {}: {}", + column_schema.name, e.what()); + } catch (const std::exception& e) { + return Status::InternalError( + "Failed to create parquet COUNT shape reader for column {}: {}", + column_schema.name, e.what()); + } case ParquetColumnSchemaKind::STRUCT: { if (column_schema.children.empty()) { return Status::NotSupported("Parquet COUNT shape reader found empty STRUCT column {}", @@ -525,7 +710,7 @@ Status ParquetColumnReaderFactory::create_count_shape_reader_impl( child_schema = column_schema.children[0].get(); } DORIS_CHECK(child_schema != nullptr); - return create_count_shape_reader_impl(*child_schema, child_projection, true, reader); + return create_count_shape_reader_impl(*child_schema, child_projection, reader); } case ParquetColumnSchemaKind::LIST: { if (column_schema.children.size() != 1) { @@ -535,7 +720,7 @@ Status ParquetColumnReaderFactory::create_count_shape_reader_impl( const auto& element_schema = *column_schema.children[0]; const auto* element_projection = format::find_child_projection(projection, element_schema.local_id); - return create_count_shape_reader_impl(element_schema, element_projection, true, reader); + return create_count_shape_reader_impl(element_schema, element_projection, reader); } case ParquetColumnSchemaKind::MAP: { if (column_schema.children.empty()) { @@ -545,7 +730,7 @@ Status ParquetColumnReaderFactory::create_count_shape_reader_impl( // The key stream defines MAP entry existence and offsets. Counting top-level MAP NULL-ness // from it avoids creating a value reader, which is the expensive path for files with huge // MAP value strings. - return create_count_shape_reader_impl(*column_schema.children[0], nullptr, true, reader); + return create_count_shape_reader_impl(*column_schema.children[0], nullptr, reader); } } return Status::NotSupported("Unsupported parquet column schema kind for COUNT column {}", diff --git a/be/src/format_v2/parquet/reader/column_reader.h b/be/src/format_v2/parquet/reader/column_reader.h index 51dbd44c11c226..d02bc6f4cc40bc 100644 --- a/be/src/format_v2/parquet/reader/column_reader.h +++ b/be/src/format_v2/parquet/reader/column_reader.h @@ -203,7 +203,6 @@ class ParquetColumnReaderFactory { std::unique_ptr* reader) const; Status create_count_shape_reader_impl(const ParquetColumnSchema& column_schema, const format::LocalColumnIndex* projection, - bool is_nested, std::unique_ptr* reader) const; Status get_record_reader(int leaf_column_id, const ::parquet::ColumnDescriptor* descriptor, diff --git a/be/src/format_v2/parquet/reader/parquet_leaf_reader.h b/be/src/format_v2/parquet/reader/parquet_leaf_reader.h index b396b35fd1f32c..089975479ac755 100644 --- a/be/src/format_v2/parquet/reader/parquet_leaf_reader.h +++ b/be/src/format_v2/parquet/reader/parquet_leaf_reader.h @@ -119,16 +119,15 @@ class ParquetLeafReader { ParquetNestedScalarBatch* batch, int16_t value_slot_repetition_level = std::numeric_limits::max()) const; - // COUNT(col) and nested-skip shape-only read path. It still calls Arrow - // RecordReader::ReadRecords() to advance the Parquet cursor and obtain def/rep levels, but - // Doris only copies levels: + // Legacy nested-skip shape-only read path. It calls Arrow RecordReader::ReadRecords() to + // advance the Parquet cursor and obtain def/rep levels, but Doris only copies levels: // - it does not build value_indices or values_column // - it does not enter DataTypeSerde::read_column_from_decoded_values() // - for Binary/FLBA, it releases and immediately discards Arrow builder chunks because that is // the RecordReader's required reset operation; it never copies them into a Doris Column - // This lets COUNT(col) on MAP/ARRAY/STRUCT evaluate top-level NULL state and lets skip advance - // nested shape without Doris-side STRING/BINARY materialization. Arrow RecordReader does not - // expose a public levels-only API, so ReadRecords may still perform required page decoding. + // This lets nested skip advance shape without Doris-side STRING/BINARY materialization. + // COUNT(col) uses CountShapeColumnReader instead, which bypasses RecordReader and therefore + // avoids Arrow binary builder materialization as well as Doris value materialization. Status read_nested_levels_batch(int64_t batch_rows, ParquetNestedScalarBatch* batch) const; private: diff --git a/be/src/format_v2/table_reader.cpp b/be/src/format_v2/table_reader.cpp index 09c5d65e82b7fd..8a19140ddf4187 100644 --- a/be/src/format_v2/table_reader.cpp +++ b/be/src/format_v2/table_reader.cpp @@ -102,6 +102,8 @@ std::string push_down_agg_to_string(TPushAggOp::type op) { return "MIX"; case TPushAggOp::COUNT_ON_INDEX: return "COUNT_ON_INDEX"; + case TPushAggOp::COUNT_NON_NULL: + return "COUNT_NON_NULL"; } return "UNKNOWN"; } @@ -472,6 +474,7 @@ Status TableReader::init(TableReadOptions&& options) { _scanner_profile = options.scanner_profile; _file_slot_descs = options.file_slot_descs; _push_down_agg_type = options.push_down_agg_type; + _count_non_null_global_index = options.count_non_null_global_index; _condition_cache_digest = options.condition_cache_digest; _projected_columns = std::move(options.projected_columns); _system_properties = create_system_properties(_scan_params); @@ -560,6 +563,7 @@ Status TableReader::_open_local_filter_exprs(const FileScanRequest& file_request bool TableReader::_should_enable_condition_cache(const FileScanRequest& file_request) const { if (_condition_cache_digest == 0 || _push_down_agg_type == TPushAggOp::type::COUNT || + _push_down_agg_type == TPushAggOp::type::COUNT_NON_NULL || _current_file_description == std::nullopt || _data_reader.reader == nullptr) { return false; } diff --git a/be/src/format_v2/table_reader.h b/be/src/format_v2/table_reader.h index b8457ee9b4a31e..a568d0f3fcfde4 100644 --- a/be/src/format_v2/table_reader.h +++ b/be/src/format_v2/table_reader.h @@ -135,6 +135,8 @@ struct TableReadOptions { const std::vector* file_slot_descs = nullptr; // Push-down aggregate type. const TPushAggOp::type push_down_agg_type = TPushAggOp::type::NONE; + // Global projected-column index targeted by COUNT_NON_NULL. + const std::optional count_non_null_global_index = std::nullopt; // Digest of stable pushed-down predicates. A zero digest disables condition cache. uint64_t condition_cache_digest = 0; }; @@ -542,12 +544,22 @@ class TableReader { bool _is_table_level_count_active() const { return _remaining_table_level_count >= 0; } - Status _materialize_count_rows(size_t rows, Block* block) const { + Status _materialize_count_rows( + size_t rows, Block* block, + std::optional non_null_global_index = std::nullopt) const { DORIS_CHECK(block != nullptr); DORIS_CHECK(block->columns() > 0 || rows == 0); + if (non_null_global_index.has_value()) { + DORIS_CHECK(non_null_global_index->value() < block->columns()); + } for (size_t column_idx = 0; column_idx < block->columns(); ++column_idx) { auto column = block->get_by_position(column_idx).type->create_column(); - column->resize(rows); + if (non_null_global_index.has_value() && column_idx == non_null_global_index->value()) { + DORIS_CHECK(column->is_nullable()); + assert_cast(*column).insert_not_null_elements(rows); + } else { + column->insert_many_defaults(rows); + } block->replace_by_position(column_idx, std::move(column)); } return Status::OK(); @@ -652,15 +664,29 @@ class TableReader { // Finalize file-local block to table/global schema block. Status finalize_chunk(Block* block, const size_t rows) { SCOPED_TIMER(_profile.finalize_timer); - size_t idx = 0; + std::vector materialized_columns; + materialized_columns.reserve(_data_reader.column_mapper->mappings().size()); for (const auto& mapping : _data_reader.column_mapper->mappings()) { ColumnPtr column; RETURN_IF_ERROR(_materialize_mapping_column(mapping, &_data_reader.block_template, rows, &column)); - block->replace_by_position(idx, IColumn::mutate(std::move(column))); - idx++; + materialized_columns.push_back(std::move(column)); } + for (size_t idx = 0; idx < materialized_columns.size(); ++idx) { + block->replace_by_position(idx, std::move(materialized_columns[idx])); + } + // Table-format virtual columns can depend on auxiliary file-local columns such as Iceberg + // row positions, so materialize them before releasing the file block. RETURN_IF_ERROR(materialize_virtual_columns(block)); + // Projection results can alias columns owned by the file-local block. Release those + // owners before requesting mutable output columns so direct mappings transfer their + // payload instead of triggering a COW deep clone. + _data_reader.block_template.clear_column_data( + cast_set(_data_reader.file_block_layout.size())); + for (size_t idx = 0; idx < block->columns(); ++idx) { + block->replace_by_position( + idx, IColumn::mutate(std::move(block->get_by_position(idx).column))); + } // Enforce CHAR/VARCHAR length declared by the table schema after all file-to-table // materialization has finished. RETURN_IF_ERROR(_truncate_char_or_varchar_columns(block)); @@ -901,6 +927,10 @@ class TableReader { block->clear_column_data(_projected_columns.size()); _aggregate_pushdown_tried = true; if (!_supports_aggregate_pushdown(_push_down_agg_type)) { + if (_push_down_agg_type == TPushAggOp::type::COUNT_NON_NULL) { + return Status::NotSupported( + "COUNT_NON_NULL requires one directly mapped file column without filters"); + } return Status::OK(); } @@ -909,6 +939,10 @@ class TableReader { FileAggregateResult file_result; const auto status = _data_reader.reader->get_aggregate_result(file_request, &file_result); if (status.is()) { + if (_push_down_agg_type == TPushAggOp::type::COUNT_NON_NULL) { + return Status::NotSupported( + "File reader does not support COUNT_NON_NULL aggregate pushdown"); + } return Status::OK(); } RETURN_IF_ERROR(status); @@ -920,8 +954,9 @@ class TableReader { } virtual bool _supports_aggregate_pushdown(TPushAggOp::type agg_type) const { - // Only COUNT and MIN/MAX can be push down. - if (agg_type != TPushAggOp::type::COUNT && agg_type != TPushAggOp::type::MINMAX) { + // Only COUNT, COUNT_NON_NULL and MIN/MAX can be pushed down. + if (agg_type != TPushAggOp::type::COUNT && agg_type != TPushAggOp::type::COUNT_NON_NULL && + agg_type != TPushAggOp::type::MINMAX) { return false; } // Aggregate pushdown returns reduced synthetic rows and may close the physical reader @@ -951,6 +986,17 @@ class TableReader { if (agg_type == TPushAggOp::type::COUNT) { return true; } + if (agg_type == TPushAggOp::type::COUNT_NON_NULL) { + DORIS_CHECK(_count_non_null_global_index.has_value()); + const auto mapping_it = std::ranges::find_if( + _data_reader.column_mapper->mappings(), [&](const auto& mapping) { + return mapping.global_index == *_count_non_null_global_index; + }); + DORIS_CHECK(mapping_it != _data_reader.column_mapper->mappings().end()); + return mapping_it->file_local_id.has_value() && + mapping_it->virtual_column_type == TableVirtualColumnType::INVALID && + mapping_it->default_expr == nullptr; + } // For MIN/MAX, only support direct file-to-table column mappings. The two emitted rows // must be enough for the upper MIN/MAX aggregate without evaluating default expressions or // virtual columns. @@ -1149,7 +1195,7 @@ class TableReader { rows, st.to_string(), mapping.debug_string()); } ColumnPtr result_column = current_block->get_by_position(res_id).column; - *column = _detach_column(std::move(result_column)); + *column = std::move(result_column); return Status::OK(); } if (mapping.default_expr != nullptr) { @@ -1159,7 +1205,7 @@ class TableReader { mapping.default_expr, current_block, &result)); ColumnPtr result_column = result.column; RETURN_IF_ERROR(_align_column_nullability(&result_column, mapping.table_type)); - *column = _detach_column(std::move(result_column)); + *column = std::move(result_column); } else { DORIS_CHECK(mapping.constant_index.has_value()); Block eval_block; @@ -1170,12 +1216,12 @@ class TableReader { mapping.default_expr, &eval_block, &result)); ColumnPtr result_column = result.column; RETURN_IF_ERROR(_align_column_nullability(&result_column, mapping.table_type)); - *column = _detach_column(std::move(result_column)); + *column = std::move(result_column); } return Status::OK(); } ColumnPtr result_column = mapping.table_type->create_column_const_with_default_value(rows); - *column = _detach_column(std::move(result_column)); + *column = std::move(result_column); return Status::OK(); } @@ -1432,25 +1478,18 @@ class TableReader { request->agg_type = agg_type; request->columns.clear(); if (agg_type == TPushAggOp::type::COUNT) { - // COUNT pushdown historically meant COUNT(*) and therefore carried no columns. For - // complex COUNT(col), materializing the full MAP/LIST/STRUCT value only to test the - // top-level NULL bit can be extremely expensive. When the scan projects exactly one - // directly-mapped complex column, pass that file column to the reader so formats such - // as Parquet can count the column shape from metadata/levels without decoding payload - // values like MAP value strings. Other COUNT cases stay on the existing row-count path - // to avoid changing count(*) semantics. - if (_data_reader.column_mapper->mappings().size() == 1) { - const auto& mapping = _data_reader.column_mapper->mappings()[0]; - if (mapping.file_local_id.has_value() && mapping.file_type != nullptr && - is_complex_type(remove_nullable(mapping.file_type)->get_primitive_type()) && - mapping.virtual_column_type == TableVirtualColumnType::INVALID && - mapping.default_expr == nullptr) { - FileAggregateRequest::Column column; - column.projection = - LocalColumnIndex::top_level(LocalColumnId(*mapping.file_local_id)); - request->columns.push_back(std::move(column)); - } - } + return Status::OK(); + } + if (agg_type == TPushAggOp::type::COUNT_NON_NULL) { + DORIS_CHECK(_count_non_null_global_index.has_value()); + const auto mapping_it = std::ranges::find_if( + _data_reader.column_mapper->mappings(), [&](const auto& mapping) { + return mapping.global_index == *_count_non_null_global_index; + }); + DORIS_CHECK(mapping_it != _data_reader.column_mapper->mappings().end()); + DORIS_CHECK(mapping_it->file_local_id.has_value()); + request->columns.push_back({.projection = LocalColumnIndex::top_level( + LocalColumnId(*mapping_it->file_local_id))}); return Status::OK(); } request->columns.reserve(_data_reader.column_mapper->mappings().size()); @@ -1469,11 +1508,15 @@ class TableReader { Status _materialize_aggregate_pushdown_rows(TPushAggOp::type agg_type, const FileAggregateResult& file_result, Block* block) { - if (agg_type == TPushAggOp::type::COUNT) { - // COUNT pushdown is not a final count value. It emits `count` default rows so the - // upper COUNT(*) aggregate can count them and produce the final result, including - // zero rows when count is 0. + if (agg_type == TPushAggOp::type::COUNT || agg_type == TPushAggOp::type::COUNT_NON_NULL) { + // Pushdown emits `count` synthetic rows rather than a final scalar. COUNT_NON_NULL + // explicitly makes its target column non-NULL, so the upper COUNT(*) or COUNT(col) + // produces the same result, including zero when no synthetic rows are emitted. DORIS_CHECK(file_result.count >= 0); + if (agg_type == TPushAggOp::type::COUNT_NON_NULL) { + return _materialize_count_rows(cast_set(file_result.count), block, + _count_non_null_global_index); + } return _materialize_count_rows(cast_set(file_result.count), block); } // MIN/MAX pushdown emits two rows, min first and max second, for each projected column. @@ -1570,6 +1613,7 @@ class TableReader { const std::vector* _file_slot_descs = nullptr; FileFormat _format; TPushAggOp::type _push_down_agg_type = TPushAggOp::type::NONE; + std::optional _count_non_null_global_index; size_t _batch_size = 0; uint64_t _condition_cache_digest = 0; segment_v2::ConditionCache::ExternalCacheKey _condition_cache_key; diff --git a/be/test/format_v2/orc/orc_reader_test.cpp b/be/test/format_v2/orc/orc_reader_test.cpp index b98e2bf14e6d0e..7a4d986613ef2d 100644 --- a/be/test/format_v2/orc/orc_reader_test.cpp +++ b/be/test/format_v2/orc/orc_reader_test.cpp @@ -4753,7 +4753,7 @@ TEST_F(NewOrcReaderTest, AggregatePushdownCountColumnUsesNonNullValueCount) { EXPECT_EQ(count_star_result.count, 400); format::FileAggregateRequest count_column_request; - count_column_request.agg_type = TPushAggOp::type::COUNT; + count_column_request.agg_type = TPushAggOp::type::COUNT_NON_NULL; count_column_request.columns.push_back({.projection = field_projection(0)}); format::FileAggregateResult count_column_result; status = reader->get_aggregate_result(count_column_request, &count_column_result); diff --git a/be/test/format_v2/parquet/parquet_reader_test.cpp b/be/test/format_v2/parquet/parquet_reader_test.cpp index 71d1cc291754ae..f70e96f41be8f9 100644 --- a/be/test/format_v2/parquet/parquet_reader_test.cpp +++ b/be/test/format_v2/parquet/parquet_reader_test.cpp @@ -1190,7 +1190,7 @@ TEST_F(NewParquetReaderTest, CountComplexColumnUsesShapeOnlyPath) { ASSERT_TRUE(reader->open(std::make_shared()).ok()); format::FileAggregateRequest request; - request.agg_type = TPushAggOp::type::COUNT; + request.agg_type = TPushAggOp::type::COUNT_NON_NULL; request.columns.push_back( {.projection = format::LocalColumnIndex::top_level(format::LocalColumnId(0))}); format::FileAggregateResult result; @@ -1199,6 +1199,8 @@ TEST_F(NewParquetReaderTest, CountComplexColumnUsesShapeOnlyPath) { // Rows are: non-empty map, NULL map, empty map, non-empty map with large value string, // non-empty map with NULL value. COUNT(arr) excludes only the top-level NULL map. EXPECT_EQ(result.count, 4); + ASSERT_NE(profile.get_counter("ArrowReadRecordsTime"), nullptr); + EXPECT_EQ(profile.get_counter("ArrowReadRecordsTime")->value(), 0); ASSERT_NE(profile.get_counter("MaterializationTime"), nullptr); EXPECT_EQ(profile.get_counter("MaterializationTime")->value(), 0); } @@ -1212,7 +1214,7 @@ TEST_F(NewParquetReaderTest, CountArrayColumnUsesLevelsOnlyPath) { ASSERT_TRUE(reader->open(std::make_shared()).ok()); format::FileAggregateRequest request; - request.agg_type = TPushAggOp::type::COUNT; + request.agg_type = TPushAggOp::type::COUNT_NON_NULL; request.columns.push_back( {.projection = format::LocalColumnIndex::top_level(format::LocalColumnId(0))}); format::FileAggregateResult result; @@ -1221,6 +1223,8 @@ TEST_F(NewParquetReaderTest, CountArrayColumnUsesLevelsOnlyPath) { // Rows are: non-empty array with a large string, NULL array, empty array, non-empty array // with NULL element, non-empty array with a large string. Only the top-level NULL is excluded. EXPECT_EQ(result.count, 4); + ASSERT_NE(profile.get_counter("ArrowReadRecordsTime"), nullptr); + EXPECT_EQ(profile.get_counter("ArrowReadRecordsTime")->value(), 0); ASSERT_NE(profile.get_counter("MaterializationTime"), nullptr); EXPECT_EQ(profile.get_counter("MaterializationTime")->value(), 0); } @@ -1234,7 +1238,7 @@ TEST_F(NewParquetReaderTest, CountStructColumnUsesLevelsOnlyPath) { ASSERT_TRUE(reader->open(std::make_shared()).ok()); format::FileAggregateRequest request; - request.agg_type = TPushAggOp::type::COUNT; + request.agg_type = TPushAggOp::type::COUNT_NON_NULL; request.columns.push_back( {.projection = format::LocalColumnIndex::top_level(format::LocalColumnId(0))}); format::FileAggregateResult result; @@ -1243,6 +1247,8 @@ TEST_F(NewParquetReaderTest, CountStructColumnUsesLevelsOnlyPath) { // The representative STRUCT leaf is the first child, a nullable STRING payload. A row with // NULL payload but non-NULL struct still counts; only the top-level NULL struct is excluded. EXPECT_EQ(result.count, 4); + ASSERT_NE(profile.get_counter("ArrowReadRecordsTime"), nullptr); + EXPECT_EQ(profile.get_counter("ArrowReadRecordsTime")->value(), 0); ASSERT_NE(profile.get_counter("MaterializationTime"), nullptr); EXPECT_EQ(profile.get_counter("MaterializationTime")->value(), 0); } @@ -1257,7 +1263,7 @@ TEST_F(NewParquetReaderTest, CountStructWithRepeatedChildUsesTopLevelRowBoundari ASSERT_TRUE(reader->open(std::make_shared()).ok()); format::FileAggregateRequest request; - request.agg_type = TPushAggOp::type::COUNT; + request.agg_type = TPushAggOp::type::COUNT_NON_NULL; request.columns.push_back({.projection = format::LocalColumnIndex::top_level( format::LocalColumnId(column_id))}); format::FileAggregateResult result; diff --git a/be/test/format_v2/table_reader_test.cpp b/be/test/format_v2/table_reader_test.cpp index 0d484f60d9e5de..3e9d9adb7971a7 100644 --- a/be/test/format_v2/table_reader_test.cpp +++ b/be/test/format_v2/table_reader_test.cpp @@ -1012,8 +1012,10 @@ struct FakeFileReaderState { bool stop_during_read = false; bool not_found_during_init = false; std::shared_ptr last_request; + std::optional last_aggregate_request; std::shared_ptr condition_cache_ctx; std::shared_ptr io_ctx; + const UInt8* string_payload = nullptr; }; class FakeFileReader final : public FileReader { @@ -1074,6 +1076,7 @@ class FakeFileReader final : public FileReader { auto column = ColumnString::create(); column->insert_data("one", 3); column->insert_data("two", 3); + _state->string_payload = column->get_chars().data(); file_block->replace_by_position(block_position.value(), make_not_null_nullable_column(std::move(column))); } else if (file_column_id == LocalColumnId(2)) { @@ -1121,10 +1124,12 @@ class FakeFileReader final : public FileReader { Status get_aggregate_result(const FileAggregateRequest& request, FileAggregateResult* result) override { DORIS_CHECK(result != nullptr); + _state->last_aggregate_request = request; if (_state->aggregate_count < 0) { return FileReader::get_aggregate_result(request, result); } - if (request.agg_type != TPushAggOp::type::COUNT) { + if (request.agg_type != TPushAggOp::type::COUNT && + request.agg_type != TPushAggOp::type::COUNT_NON_NULL) { return Status::NotSupported("fake reader only supports COUNT aggregate pushdown"); } if (_state->stop_during_aggregate) { @@ -1671,6 +1676,42 @@ TEST(TableReaderTest, SlotlessConjunctDisablesAggregatePushdown) { ASSERT_TRUE(reader.close().ok()); } +TEST(TableReaderTest, DirectMappingTransfersFileColumnOwnership) { + std::vector file_schema; + file_schema.push_back(make_file_column(1, "value", std::make_shared())); + + std::vector projected_columns; + projected_columns.push_back(make_table_column(1, "value", std::make_shared())); + set_name_identifiers(&projected_columns); + + RuntimeState state {TQueryOptions(), TQueryGlobals()}; + auto fake_state = std::make_shared(); + FakeTableReader reader(file_schema, fake_state); + ASSERT_TRUE(reader.init({ + .projected_columns = projected_columns, + .conjuncts = {}, + .format = FileFormat::PARQUET, + .scan_params = nullptr, + .io_ctx = nullptr, + .runtime_state = &state, + .scanner_profile = nullptr, + }) + .ok()); + + SplitReadOptions split_options; + split_options.current_range.__set_path("fake-table-reader-input"); + ASSERT_TRUE(reader.prepare_split(split_options).ok()); + + Block block = build_table_block(projected_columns); + bool eos = false; + ASSERT_TRUE(reader.get_block(&block, &eos).ok()); + ASSERT_FALSE(eos); + ASSERT_NE(fake_state->string_payload, nullptr); + + const auto& result = assert_cast(expect_not_null_table_column(block, 0)); + EXPECT_EQ(result.get_chars().data(), fake_state->string_payload); +} + TEST(TableReaderTest, AbortSplitClearsReaderAfterIgnorableNotFound) { std::vector file_schema; file_schema.push_back(make_file_column(0, "id", std::make_shared())); @@ -1755,6 +1796,51 @@ TEST(TableReaderTest, PushDownCountRecordsReaderRowsBeforeClosingReader) { EXPECT_EQ(fake_state->close_count, 1); } +TEST(TableReaderTest, CountNonNullUsesExplicitTargetColumn) { + std::vector file_schema; + file_schema.push_back(make_file_column(0, "id", std::make_shared())); + file_schema.push_back( + make_file_column(1, "value", make_nullable(std::make_shared()))); + + std::vector projected_columns; + projected_columns.push_back(make_table_column(0, "id", std::make_shared())); + projected_columns.push_back( + make_table_column(1, "value", make_nullable(std::make_shared()))); + set_name_identifiers(&projected_columns); + + RuntimeState state {TQueryOptions(), TQueryGlobals()}; + auto fake_state = std::make_shared(); + fake_state->aggregate_count = 3; + FakeTableReader reader(file_schema, fake_state); + ASSERT_TRUE(reader.init({ + .projected_columns = projected_columns, + .conjuncts = {}, + .format = FileFormat::PARQUET, + .scan_params = nullptr, + .io_ctx = nullptr, + .runtime_state = &state, + .scanner_profile = nullptr, + .push_down_agg_type = TPushAggOp::type::COUNT_NON_NULL, + .count_non_null_global_index = GlobalIndex(1), + }) + .ok()); + + SplitReadOptions split_options; + split_options.current_range.__set_path("fake-table-reader-input"); + ASSERT_TRUE(reader.prepare_split(split_options).ok()); + + Block block = build_table_block(projected_columns); + bool eos = false; + ASSERT_TRUE(reader.get_block(&block, &eos).ok()); + ASSERT_FALSE(eos); + ASSERT_TRUE(fake_state->last_aggregate_request.has_value()); + EXPECT_EQ(fake_state->last_aggregate_request->agg_type, TPushAggOp::type::COUNT_NON_NULL); + ASSERT_EQ(fake_state->last_aggregate_request->columns.size(), 1); + EXPECT_EQ(fake_state->last_aggregate_request->columns[0].projection.local_id(), 1); + EXPECT_EQ(block.rows(), 3); + expect_not_null_table_column(block, 1); +} + TEST(TableReaderTest, PushDownCountStopConvertsAggregateEndOfFileToEos) { std::vector file_schema; file_schema.push_back(make_file_column(0, "id", std::make_shared())); @@ -1873,10 +1959,11 @@ TEST(TableReaderTest, DebugStringCoversReaderStateAndEnumNames) { std::string::npos); } - const std::vector agg_ops {TPushAggOp::type::NONE, TPushAggOp::type::MINMAX, - TPushAggOp::type::MIX, - TPushAggOp::type::COUNT_ON_INDEX}; - const std::vector agg_names {"NONE", "MINMAX", "MIX", "COUNT_ON_INDEX"}; + const std::vector agg_ops { + TPushAggOp::type::NONE, TPushAggOp::type::MINMAX, TPushAggOp::type::MIX, + TPushAggOp::type::COUNT_ON_INDEX, TPushAggOp::type::COUNT_NON_NULL}; + const std::vector agg_names {"NONE", "MINMAX", "MIX", "COUNT_ON_INDEX", + "COUNT_NON_NULL"}; for (size_t idx = 0; idx < agg_ops.size(); ++idx) { TableReader enum_reader; ASSERT_TRUE(enum_reader diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/FileScanNode.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/FileScanNode.java index 466dad11d5dc49..585182c0bf1067 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/FileScanNode.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/FileScanNode.java @@ -95,6 +95,9 @@ protected void toThrift(TPlanNode planNode) { if (desc.getTable() != null) { fileScanNode.setTableName(desc.getTable().getName()); } + if (countNonNullSlotId != null) { + fileScanNode.setCountNonNullSlotId(countNonNullSlotId.asInt()); + } planNode.setFileScanNode(fileScanNode); super.toThrift(planNode); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java index 01d9139577578b..1c87cb9f94fa62 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java @@ -844,12 +844,27 @@ public PlanFragment visitPhysicalHudiScan(PhysicalHudiScan hudiScan, PlanTransla return getPlanFragmentForPhysicalFileScan(hudiScan, context, hudiScanNode); } + private void setFilePushDownAggregate(PhysicalRelation relation, ScanNode scanNode, + PlanTranslatorContext context) { + TPushAggOp pushAggOp = context.getRelationPushAggOp(relation.getRelationId()); + scanNode.setPushDownAggNoGrouping(pushAggOp); + if (pushAggOp == TPushAggOp.COUNT_NON_NULL) { + Optional targetSlot = context.getRelationPushAggSlot(relation.getRelationId()); + Preconditions.checkState(targetSlot.isPresent(), + "COUNT_NON_NULL target is missing for relation " + relation.getRelationId()); + SlotRef slotRef = context.findSlotRef(targetSlot.get()); + Preconditions.checkState(slotRef != null, + "COUNT_NON_NULL target slot is missing for relation " + relation.getRelationId()); + scanNode.setCountNonNullSlotId(slotRef.getSlotId()); + } + } + @NotNull private PlanFragment getPlanFragmentForPhysicalFileScan(PhysicalFileScan fileScan, PlanTranslatorContext context, ScanNode scanNode) { scanNode.setNereidsId(fileScan.getId()); context.getNereidsIdToPlanNodeIdMap().put(fileScan.getId(), scanNode.getId()); - scanNode.setPushDownAggNoGrouping(context.getRelationPushAggOp(fileScan.getRelationId())); + setFilePushDownAggregate(fileScan, scanNode, context); scanNode.setHasPartitionPredicate(fileScan.hasPartitionPredicate()); if (fileScan.getStats() != null) { @@ -1105,6 +1120,7 @@ public PlanFragment visitPhysicalTVFRelation(PhysicalTVFRelation tvfRelation, Pl TableValuedFunctionIf catalogFunction = tvfRelation.getFunction().getCatalogFunction(); SessionVariable sv = ConnectContext.get().getSessionVariable(); ScanNode scanNode = catalogFunction.getScanNode(context.nextPlanNodeId(), tupleDescriptor, sv); + setFilePushDownAggregate(tvfRelation, scanNode, context); scanNode.setDistributeExprLists(getDistributeExpr(tvfRelation)); scanNode.setNereidsId(tvfRelation.getId()); context.getNereidsIdToPlanNodeIdMap().put(tvfRelation.getId(), scanNode.getId()); @@ -1370,8 +1386,9 @@ public PlanFragment visitPhysicalBucketedHashAggregate( public PlanFragment visitPhysicalStorageLayerAggregate( PhysicalStorageLayerAggregate storageLayerAggregate, PlanTranslatorContext context) { Preconditions.checkState((storageLayerAggregate.getRelation() instanceof PhysicalOlapScan - || storageLayerAggregate.getRelation() instanceof PhysicalFileScan), - "PhysicalStorageLayerAggregate only support PhysicalOlapScan and PhysicalFileScan: " + || storageLayerAggregate.getRelation() instanceof PhysicalFileScan + || storageLayerAggregate.getRelation() instanceof PhysicalTVFRelation), + "PhysicalStorageLayerAggregate only supports file and OLAP scans: " + storageLayerAggregate.getRelation().getClass().getName()); TPushAggOp pushAggOp; @@ -1379,6 +1396,9 @@ public PlanFragment visitPhysicalStorageLayerAggregate( case COUNT: pushAggOp = TPushAggOp.COUNT; break; + case COUNT_NON_NULL: + pushAggOp = TPushAggOp.COUNT_NON_NULL; + break; case COUNT_ON_MATCH: pushAggOp = TPushAggOp.COUNT_ON_INDEX; break; @@ -1393,14 +1413,21 @@ public PlanFragment visitPhysicalStorageLayerAggregate( + storageLayerAggregate.getAggOp()); } - if (storageLayerAggregate.getRelation() instanceof PhysicalFileScan - && pushAggOp.equals(TPushAggOp.COUNT) + if ((storageLayerAggregate.getRelation() instanceof PhysicalFileScan + || storageLayerAggregate.getRelation() instanceof PhysicalTVFRelation) + && (pushAggOp.equals(TPushAggOp.COUNT) || pushAggOp.equals(TPushAggOp.COUNT_NON_NULL)) && !ConnectContext.get().getSessionVariable().isEnableCountPushDownForExternalTable()) { pushAggOp = TPushAggOp.NONE; } context.setRelationPushAggOp( storageLayerAggregate.getRelation().getRelationId(), pushAggOp); + if (pushAggOp == TPushAggOp.COUNT_NON_NULL) { + Preconditions.checkState(storageLayerAggregate.getAggSlot().isPresent(), + "COUNT_NON_NULL requires one target slot"); + context.setRelationPushAggSlot(storageLayerAggregate.getRelation().getRelationId(), + storageLayerAggregate.getAggSlot().get()); + } PlanFragment planFragment = storageLayerAggregate.getRelation().accept(this, context); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PlanTranslatorContext.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PlanTranslatorContext.java index f680936b5e5ef4..f2b00cf14086de 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PlanTranslatorContext.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PlanTranslatorContext.java @@ -106,6 +106,7 @@ public class PlanTranslatorContext { private final Map cteScanNodeMap = Maps.newHashMap(); private final Map tablePushAggOp = Maps.newHashMap(); + private final Map tablePushAggSlot = Maps.newHashMap(); private final Map> statsUnknownColumnsMap = Maps.newHashMap(); @@ -430,6 +431,14 @@ public TPushAggOp getRelationPushAggOp(RelationId relationId) { return tablePushAggOp.getOrDefault(relationId, TPushAggOp.NONE); } + public void setRelationPushAggSlot(RelationId relationId, ExprId exprId) { + tablePushAggSlot.put(relationId, exprId); + } + + public Optional getRelationPushAggSlot(RelationId relationId) { + return Optional.ofNullable(tablePushAggSlot.get(relationId)); + } + public boolean isTopMaterializeNode() { return isTopMaterializeNode; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java index 1e4f38fd005c80..4a810e9087482a 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java @@ -575,6 +575,8 @@ public enum RuleType { STORAGE_LAYER_AGGREGATE_WITH_PROJECT(RuleTypeClass.IMPLEMENTATION), STORAGE_LAYER_AGGREGATE_WITHOUT_PROJECT_FOR_FILE_SCAN(RuleTypeClass.IMPLEMENTATION), STORAGE_LAYER_AGGREGATE_WITH_PROJECT_FOR_FILE_SCAN(RuleTypeClass.IMPLEMENTATION), + STORAGE_LAYER_AGGREGATE_WITHOUT_PROJECT_FOR_TVF(RuleTypeClass.IMPLEMENTATION), + STORAGE_LAYER_AGGREGATE_WITH_PROJECT_FOR_TVF(RuleTypeClass.IMPLEMENTATION), STORAGE_LAYER_WITH_PROJECT_NO_SLOT_REF(RuleTypeClass.IMPLEMENTATION), STORAGE_LAYER_AGGREGATE_MINMAX_ON_UNIQUE(RuleTypeClass.IMPLEMENTATION), STORAGE_LAYER_AGGREGATE_MINMAX_ON_UNIQUE_WITHOUT_PROJECT(RuleTypeClass.IMPLEMENTATION), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/AggregateStrategies.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/AggregateStrategies.java index 975cfaf973c408..b5ca42362ce2f4 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/AggregateStrategies.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/AggregateStrategies.java @@ -33,6 +33,7 @@ import org.apache.doris.nereids.rules.expression.rules.FoldConstantRuleOnFE; import org.apache.doris.nereids.trees.expressions.Alias; import org.apache.doris.nereids.trees.expressions.Cast; +import org.apache.doris.nereids.trees.expressions.ExprId; import org.apache.doris.nereids.trees.expressions.Expression; import org.apache.doris.nereids.trees.expressions.IsNull; import org.apache.doris.nereids.trees.expressions.Or; @@ -53,12 +54,16 @@ import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; import org.apache.doris.nereids.trees.plans.logical.LogicalProject; import org.apache.doris.nereids.trees.plans.logical.LogicalRelation; +import org.apache.doris.nereids.trees.plans.logical.LogicalTVFRelation; +import org.apache.doris.nereids.trees.plans.logical.SupportPruneNestedColumn; import org.apache.doris.nereids.trees.plans.physical.PhysicalFileScan; import org.apache.doris.nereids.trees.plans.physical.PhysicalOlapScan; import org.apache.doris.nereids.trees.plans.physical.PhysicalStorageLayerAggregate; import org.apache.doris.nereids.trees.plans.physical.PhysicalStorageLayerAggregate.PushDownAggOp; +import org.apache.doris.nereids.trees.plans.physical.PhysicalTVFRelation; import org.apache.doris.nereids.util.ExpressionUtils; import org.apache.doris.qe.ConnectContext; +import org.apache.doris.tablefunction.ExternalFileTableValuedFunction; import com.google.common.collect.ImmutableList; @@ -258,6 +263,26 @@ public List buildRules() { LogicalFileScan fileScan = project.child(); return storageLayerAggregate(agg, project, fileScan, ctx.cascadesContext); }) + ), + RuleType.STORAGE_LAYER_AGGREGATE_WITHOUT_PROJECT_FOR_TVF.build( + logicalAggregate( + logicalTVFRelation() + ) + .when(agg -> agg.isNormalized() && enablePushDownNoGroupAgg()) + .thenApply(ctx -> storageLayerAggregate(ctx.root, null, ctx.root.child(), ctx.cascadesContext)) + ), + RuleType.STORAGE_LAYER_AGGREGATE_WITH_PROJECT_FOR_TVF.build( + logicalAggregate( + logicalProject( + logicalTVFRelation() + ) + ).when(agg -> agg.isNormalized() && enablePushDownNoGroupAgg()) + .thenApply(ctx -> { + LogicalAggregate> agg = ctx.root; + LogicalProject project = agg.child(); + LogicalTVFRelation tvf = project.child(); + return storageLayerAggregate(agg, project, tvf, ctx.cascadesContext); + }) ) ); } @@ -549,7 +574,20 @@ private LogicalAggregate storageLayerAggregate( LogicalRelation logicalScan, CascadesContext cascadesContext) { final LogicalAggregate canNotPush = aggregate; - if (!(logicalScan instanceof LogicalOlapScan) && !(logicalScan instanceof LogicalFileScan)) { + if (!(logicalScan instanceof LogicalOlapScan) && !(logicalScan instanceof LogicalFileScan) + && !(logicalScan instanceof LogicalTVFRelation)) { + return canNotPush; + } + if (logicalScan instanceof LogicalTVFRelation + && !(((LogicalTVFRelation) logicalScan).getFunction().getCatalogFunction() + instanceof ExternalFileTableValuedFunction)) { + return canNotPush; + } + if (logicalScan instanceof LogicalTVFRelation + && (!(((LogicalTVFRelation) logicalScan).getFunction() + instanceof SupportPruneNestedColumn) + || !((SupportPruneNestedColumn) ((LogicalTVFRelation) logicalScan) + .getFunction()).supportPruneNestedColumn())) { return canNotPush; } @@ -691,7 +729,24 @@ private LogicalAggregate storageLayerAggregate( List usedSlotInTable = (List) Project.findProject(aggUsedSlots, logicalScan.getOutput()); + Optional pushDownAggSlot = Optional.empty(); + for (SlotReference slot : usedSlotInTable) { + if (logicalScan instanceof LogicalTVFRelation) { + if (mergeOp == PushDownAggOp.COUNT + && slot.nullable() + && checkNullSlots.contains(slot)) { + boolean canPushNullableCount = aggregateFunctions.size() == 1 + && usedSlotInTable.size() == 1 + && enableFileScannerV2(); + if (!canPushNullableCount) { + return canNotPush; + } + mergeOp = PushDownAggOp.COUNT_NON_NULL; + pushDownAggSlot = Optional.of(slot.getExprId()); + } + continue; + } Optional optionalColumn = slot.getOriginalColumn(); if (!optionalColumn.isPresent()) { // virtual columns (e.g., generated from MATCH_ALL expressions) do not have @@ -729,6 +784,12 @@ private LogicalAggregate storageLayerAggregate( } } + if (logicalScan instanceof LogicalTVFRelation + && mergeOp != PushDownAggOp.COUNT + && mergeOp != PushDownAggOp.COUNT_NON_NULL) { + return canNotPush; + } + if (logicalScan instanceof LogicalOlapScan) { PhysicalOlapScan physicalScan = (PhysicalOlapScan) new LogicalOlapScanToPhysicalOlapScan() .build() @@ -738,11 +799,12 @@ private LogicalAggregate storageLayerAggregate( if (project != null) { return aggregate.withChildren(ImmutableList.of( project.withChildren( - ImmutableList.of(new PhysicalStorageLayerAggregate(physicalScan, mergeOp))) + ImmutableList.of(new PhysicalStorageLayerAggregate( + physicalScan, mergeOp, pushDownAggSlot))) )); } else { return aggregate.withChildren(ImmutableList.of( - new PhysicalStorageLayerAggregate(physicalScan, mergeOp) + new PhysicalStorageLayerAggregate(physicalScan, mergeOp, pushDownAggSlot) )); } @@ -754,11 +816,30 @@ private LogicalAggregate storageLayerAggregate( if (project != null) { return aggregate.withChildren(ImmutableList.of( project.withChildren( - ImmutableList.of(new PhysicalStorageLayerAggregate(physicalScan, mergeOp))) + ImmutableList.of(new PhysicalStorageLayerAggregate( + physicalScan, mergeOp, pushDownAggSlot))) + )); + } else { + return aggregate.withChildren(ImmutableList.of( + new PhysicalStorageLayerAggregate(physicalScan, mergeOp, pushDownAggSlot) + )); + } + + } else if (logicalScan instanceof LogicalTVFRelation) { + PhysicalTVFRelation physicalScan = + (PhysicalTVFRelation) new LogicalTVFRelationToPhysicalTVFRelation() + .build() + .transform(logicalScan, cascadesContext) + .get(0); + if (project != null) { + return aggregate.withChildren(ImmutableList.of( + project.withChildren( + ImmutableList.of(new PhysicalStorageLayerAggregate( + physicalScan, mergeOp, pushDownAggSlot))) )); } else { return aggregate.withChildren(ImmutableList.of( - new PhysicalStorageLayerAggregate(physicalScan, mergeOp) + new PhysicalStorageLayerAggregate(physicalScan, mergeOp, pushDownAggSlot) )); } @@ -772,6 +853,11 @@ private boolean enablePushDownStringMinMax() { return connectContext != null && connectContext.getSessionVariable().isEnablePushDownStringMinMax(); } + private boolean enableFileScannerV2() { + ConnectContext connectContext = ConnectContext.get(); + return connectContext != null && connectContext.getSessionVariable().enableFileScannerV2; + } + private boolean enablePushDownNoGroupAgg() { ConnectContext connectContext = ConnectContext.get(); return connectContext == null || connectContext.getSessionVariable().enablePushDownNoGroupAgg(); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalStorageLayerAggregate.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalStorageLayerAggregate.java index 2db9c1afc27492..7aa911ef16d4a6 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalStorageLayerAggregate.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalStorageLayerAggregate.java @@ -18,8 +18,11 @@ package org.apache.doris.nereids.trees.plans.physical; import org.apache.doris.nereids.memo.GroupExpression; +import org.apache.doris.nereids.properties.DataTrait; import org.apache.doris.nereids.properties.LogicalProperties; import org.apache.doris.nereids.properties.PhysicalProperties; +import org.apache.doris.nereids.trees.expressions.ExprId; +import org.apache.doris.nereids.trees.expressions.Slot; import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction; import org.apache.doris.nereids.trees.expressions.functions.agg.Count; import org.apache.doris.nereids.trees.expressions.functions.agg.Max; @@ -30,7 +33,6 @@ import org.apache.doris.nereids.util.Utils; import org.apache.doris.statistics.Statistics; -import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import java.util.List; @@ -39,25 +41,33 @@ import java.util.Optional; /** PhysicalStorageLayerAggregate */ -public class PhysicalStorageLayerAggregate extends PhysicalCatalogRelation { +public class PhysicalStorageLayerAggregate extends PhysicalRelation { - private final PhysicalCatalogRelation relation; + private final PhysicalRelation relation; private final PushDownAggOp aggOp; + private final Optional aggSlot; - public PhysicalStorageLayerAggregate(PhysicalCatalogRelation relation, PushDownAggOp aggOp) { - super(relation.getRelationId(), relation.getType(), relation.getTable(), relation.getQualifier(), - Optional.empty(), relation.getLogicalProperties(), ImmutableList.of()); + public PhysicalStorageLayerAggregate(PhysicalRelation relation, PushDownAggOp aggOp) { + this(relation, aggOp, Optional.empty()); + } + + public PhysicalStorageLayerAggregate(PhysicalRelation relation, PushDownAggOp aggOp, + Optional aggSlot) { + super(relation.getRelationId(), relation.getType(), Optional.empty(), relation.getLogicalProperties()); this.relation = Objects.requireNonNull(relation, "relation cannot be null"); this.aggOp = Objects.requireNonNull(aggOp, "aggOp cannot be null"); + this.aggSlot = Objects.requireNonNull(aggSlot, "aggSlot cannot be null"); } - public PhysicalStorageLayerAggregate(PhysicalCatalogRelation relation, PushDownAggOp aggOp, + public PhysicalStorageLayerAggregate(PhysicalRelation relation, PushDownAggOp aggOp, + Optional aggSlot, Optional groupExpression, LogicalProperties logicalProperties, PhysicalProperties physicalProperties, Statistics statistics) { - super(relation.getRelationId(), relation.getType(), relation.getTable(), relation.getQualifier(), - groupExpression, logicalProperties, physicalProperties, statistics, ImmutableList.of()); + super(relation.getRelationId(), relation.getType(), groupExpression, logicalProperties, + physicalProperties, statistics); this.relation = Objects.requireNonNull(relation, "relation cannot be null"); this.aggOp = Objects.requireNonNull(aggOp, "aggOp cannot be null"); + this.aggSlot = Objects.requireNonNull(aggSlot, "aggSlot cannot be null"); } public PhysicalRelation getRelation() { @@ -68,6 +78,45 @@ public PushDownAggOp getAggOp() { return aggOp; } + public Optional getAggSlot() { + return aggSlot; + } + + @Override + public DataTrait computeDataTrait() { + return relation.getLogicalProperties().getTrait(); + } + + @Override + public void computeUnique(DataTrait.Builder builder) { + builder.addUniqueSlot(relation.getLogicalProperties().getTrait()); + } + + @Override + public void computeUniform(DataTrait.Builder builder) { + builder.addUniformSlot(relation.getLogicalProperties().getTrait()); + } + + @Override + public void computeEqualSet(DataTrait.Builder builder) { + builder.addEqualSet(relation.getLogicalProperties().getTrait()); + } + + @Override + public void computeFd(DataTrait.Builder builder) { + builder.addFuncDepsDG(relation.getLogicalProperties().getTrait()); + } + + @Override + public List computeOutput() { + return relation.getOutput(); + } + + @Override + public boolean canPushDownRuntimeFilter() { + return relation.canPushDownRuntimeFilter(); + } + @Override public R accept(PlanVisitor visitor, C context) { return visitor.visitPhysicalStorageLayerAggregate(this, context); @@ -82,35 +131,52 @@ public String toString() { ); } + @Override + public String shapeInfo() { + if (!(relation instanceof PhysicalCatalogRelation)) { + return super.shapeInfo(); + } + PhysicalCatalogRelation catalogRelation = (PhysicalCatalogRelation) relation; + StringBuilder shapeBuilder = new StringBuilder(getClass().getSimpleName()) + .append("[").append(catalogRelation.getTable().getName()).append("]"); + if (!getAppliedRuntimeFilters().isEmpty()) { + shapeBuilder.append(" apply RFs:"); + getAppliedRuntimeFilters().forEach( + runtimeFilter -> shapeBuilder.append(" RF").append(runtimeFilter.getId().asInt())); + } + return shapeBuilder.toString(); + } + public PhysicalStorageLayerAggregate withPhysicalOlapScan(PhysicalOlapScan physicalOlapScan) { - return AbstractPlan.copyWithSameId(this, () -> new PhysicalStorageLayerAggregate(physicalOlapScan, aggOp)); + return AbstractPlan.copyWithSameId(this, + () -> new PhysicalStorageLayerAggregate(physicalOlapScan, aggOp, aggSlot)); } @Override public PhysicalStorageLayerAggregate withGroupExpression(Optional groupExpression) { return AbstractPlan.copyWithSameId(this, () -> new PhysicalStorageLayerAggregate(relation, aggOp, - groupExpression, getLogicalProperties(), physicalProperties, statistics)); + aggSlot, groupExpression, getLogicalProperties(), physicalProperties, statistics)); } @Override public Plan withGroupExprLogicalPropChildren(Optional groupExpression, Optional logicalProperties, List children) { return AbstractPlan.copyWithSameId(this, () -> new PhysicalStorageLayerAggregate(relation, aggOp, - groupExpression, logicalProperties.get(), physicalProperties, statistics)); + aggSlot, groupExpression, logicalProperties.get(), physicalProperties, statistics)); } @Override public PhysicalPlan withPhysicalPropertiesAndStats(PhysicalProperties physicalProperties, Statistics statistics) { return AbstractPlan.copyWithSameId(this, () -> new PhysicalStorageLayerAggregate( - (PhysicalCatalogRelation) relation.withPhysicalPropertiesAndStats(null, statistics), - aggOp, groupExpression, + (PhysicalRelation) relation.withPhysicalPropertiesAndStats(null, statistics), + aggOp, aggSlot, groupExpression, getLogicalProperties(), physicalProperties, statistics)); } /** PushAggOp */ public enum PushDownAggOp { - COUNT, MIN_MAX, MIX, COUNT_ON_MATCH; + COUNT, COUNT_NON_NULL, MIN_MAX, MIX, COUNT_ON_MATCH; /** supportedFunctions */ public static Map, PushDownAggOp> supportedFunctions() { diff --git a/fe/fe-core/src/main/java/org/apache/doris/planner/PlanNode.java b/fe/fe-core/src/main/java/org/apache/doris/planner/PlanNode.java index c1f97fdeb96e20..eb8df621476390 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/planner/PlanNode.java +++ b/fe/fe-core/src/main/java/org/apache/doris/planner/PlanNode.java @@ -771,11 +771,16 @@ public void setCardinalityAfterFilter(long cardinalityAfterFilter) { } protected TPushAggOp pushDownAggNoGroupingOp = TPushAggOp.NONE; + protected SlotId countNonNullSlotId; public void setPushDownAggNoGrouping(TPushAggOp pushDownAggNoGroupingOp) { this.pushDownAggNoGroupingOp = pushDownAggNoGroupingOp; } + public void setCountNonNullSlotId(SlotId countNonNullSlotId) { + this.countNonNullSlotId = countNonNullSlotId; + } + public void setChildrenDistributeExprLists(List> childrenDistributeExprLists) { this.childrenDistributeExprLists = childrenDistributeExprLists; } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PhysicalStorageLayerAggregateTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PhysicalStorageLayerAggregateTest.java index bf5b74f2ea6f25..4227b8e1f8a4a9 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PhysicalStorageLayerAggregateTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PhysicalStorageLayerAggregateTest.java @@ -17,27 +17,45 @@ package org.apache.doris.nereids.rules.rewrite; +import org.apache.doris.catalog.ArrayType; +import org.apache.doris.catalog.Column; +import org.apache.doris.catalog.FunctionGenTable; +import org.apache.doris.catalog.TableIf; +import org.apache.doris.catalog.Type; +import org.apache.doris.common.AnalysisException; import org.apache.doris.nereids.CascadesContext; import org.apache.doris.nereids.rules.Rule; import org.apache.doris.nereids.rules.RulePromise; import org.apache.doris.nereids.rules.RuleType; import org.apache.doris.nereids.rules.implementation.AggregateStrategies; import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.Properties; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator; import org.apache.doris.nereids.trees.expressions.functions.agg.Count; import org.apache.doris.nereids.trees.expressions.functions.agg.Max; import org.apache.doris.nereids.trees.expressions.functions.agg.Min; import org.apache.doris.nereids.trees.expressions.functions.scalar.Ln; +import org.apache.doris.nereids.trees.expressions.functions.table.Hdfs; +import org.apache.doris.nereids.trees.plans.RelationId; import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.trees.plans.logical.LogicalTVFRelation; import org.apache.doris.nereids.trees.plans.physical.PhysicalStorageLayerAggregate.PushDownAggOp; +import org.apache.doris.nereids.trees.plans.physical.PhysicalTVFRelation; +import org.apache.doris.nereids.types.DataType; import org.apache.doris.nereids.util.MemoPatternMatchSupported; import org.apache.doris.nereids.util.MemoTestUtils; import org.apache.doris.nereids.util.PlanChecker; import org.apache.doris.nereids.util.PlanConstructor; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.tablefunction.ExternalFileTableValuedFunction; +import org.apache.doris.tablefunction.TableValuedFunctionIf; import com.google.common.collect.ImmutableList; import org.junit.jupiter.api.Test; +import org.mockito.Mockito; import java.util.Collections; import java.util.Optional; @@ -75,7 +93,8 @@ public void testWithoutProject() { .applyImplementation(storageLayerAggregateWithoutProject()) .matches( logicalAggregate( - physicalStorageLayerAggregate().when(agg -> agg.getAggOp() == PushDownAggOp.COUNT) + physicalStorageLayerAggregate().when(agg -> agg.getAggOp() == PushDownAggOp.COUNT + && agg.shapeInfo().equals("PhysicalStorageLayerAggregate[tbl]")) ) ); @@ -189,6 +208,42 @@ void testProjectionCheck() { ); } + @Test + void testNullableTvfCountWithoutOriginalColumnPushesTargetColumnShape() throws AnalysisException { + ExternalFileTableValuedFunction catalogFunction = + Mockito.mock(ExternalFileTableValuedFunction.class); + Column arrayColumn = new Column("arr", ArrayType.create(Type.STRING, true), true); + FunctionGenTable table = new FunctionGenTable(1, "hdfs", + TableIf.TableType.TABLE_VALUED_FUNCTION, ImmutableList.of(arrayColumn), catalogFunction); + Hdfs function = new Hdfs( + new Properties(Collections.singletonMap("format", "parquet"))) { + @Override + protected TableValuedFunctionIf toCatalogFunction() { + return catalogFunction; + } + }; + Mockito.when(catalogFunction.getTable()).thenReturn(table); + + SlotReference arraySlot = new SlotReference(StatementScopeIdGenerator.newExprId(), "arr", + DataType.fromCatalogType(arrayColumn.getType()), true, ImmutableList.of("hdfs")); + LogicalTVFRelation tvf = new LogicalTVFRelation( + new RelationId(1), function, ImmutableList.of()).withCachedOutputs(ImmutableList.of(arraySlot)); + LogicalAggregate aggregate = new LogicalAggregate<>( + Collections.emptyList(), ImmutableList.of(new Alias(new Count(arraySlot), "count")), + true, Optional.empty(), tvf); + + ConnectContext connectContext = MemoTestUtils.createConnectContext(); + connectContext.getSessionVariable().enableFileScannerV2 = true; + CascadesContext context = MemoTestUtils.createCascadesContext(connectContext, aggregate); + + PlanChecker.from(context) + .applyImplementation(storageLayerAggregateWithoutProjectForTvf()) + .matches(logicalAggregate(physicalStorageLayerAggregate().when(agg -> + agg.getRelation() instanceof PhysicalTVFRelation + && agg.getAggOp() == PushDownAggOp.COUNT_NON_NULL + && agg.getAggSlot().equals(Optional.of(arraySlot.getExprId()))))); + } + private Rule storageLayerAggregateWithoutProject() { return new AggregateStrategies().buildRules() .stream() @@ -204,4 +259,13 @@ private Rule storageLayerAggregateWithProject() { .findFirst() .get(); } + + private Rule storageLayerAggregateWithoutProjectForTvf() { + return new AggregateStrategies().buildRules() + .stream() + .filter(rule -> rule.getRuleType() + == RuleType.STORAGE_LAYER_AGGREGATE_WITHOUT_PROJECT_FOR_TVF) + .findFirst() + .get(); + } } diff --git a/gensrc/thrift/PlanNodes.thrift b/gensrc/thrift/PlanNodes.thrift index c2dd42bd522a16..e57ec536db963b 100644 --- a/gensrc/thrift/PlanNodes.thrift +++ b/gensrc/thrift/PlanNodes.thrift @@ -836,6 +836,8 @@ struct TBrokerScanNode { struct TFileScanNode { 1: optional Types.TTupleId tuple_id 2: optional string table_name + // Target slot for COUNT_NON_NULL. COUNT without this field keeps COUNT(*) semantics. + 3: optional Types.TSlotId count_non_null_slot_id } struct TRecCTETarget { @@ -971,7 +973,8 @@ enum TPushAggOp { MINMAX = 1, COUNT = 2, MIX = 3, - COUNT_ON_INDEX = 4 + COUNT_ON_INDEX = 4, + COUNT_NON_NULL = 5 } struct TScoreRangeInfo { diff --git a/regression-test/suites/external_table_p0/tvf/test_hdfs_parquet_group0.groovy b/regression-test/suites/external_table_p0/tvf/test_hdfs_parquet_group0.groovy index 531bc0deac22b1..daad72f531ecb5 100644 --- a/regression-test/suites/external_table_p0/tvf/test_hdfs_parquet_group0.groovy +++ b/regression-test/suites/external_table_p0/tvf/test_hdfs_parquet_group0.groovy @@ -16,6 +16,8 @@ // under the License. suite("test_hdfs_parquet_group0", "p0,external") { + sql "set enable_file_scanner_v2 = true" + sql "set enable_count_push_down_for_external_table = true" String hdfs_port = context.config.otherConfigs.get("hive2HdfsPort") String externalEnvIp = context.config.otherConfigs.get("externalEnvIp") @@ -104,6 +106,13 @@ suite("test_hdfs_parquet_group0", "p0,external") { uri = "${defaultFS}" + "/user/doris/tvf_data/test_hdfs_parquet/group0/large_string_map.brotli.parquet" + explain { + sql """ select count(arr) from HDFS( + "uri" = "${uri}", + "hadoop.username" = "${hdfsUserName}", + "format" = "parquet"); """ + contains "pushdown agg=COUNT_NON_NULL" + } order_qt_test_11 """ select count(arr) from HDFS( "uri" = "${uri}", "hadoop.username" = "${hdfsUserName}",