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}",