diff --git a/src/Functions/geohashesInBox.cpp b/src/Functions/geohashesInBox.cpp index 7bb02f45e91..ebd476c650c 100644 --- a/src/Functions/geohashesInBox.cpp +++ b/src/Functions/geohashesInBox.cpp @@ -59,24 +59,39 @@ public: bool useDefaultImplementationForConstants() const override { return true; } template - void execute(const IColumn * lon_min_column, - const IColumn * lat_min_column, - const IColumn * lon_max_column, - const IColumn * lat_max_column, - const IColumn * precision_column, - ColumnPtr & result) const + void execute( + const IColumn * lon_min_column, + const IColumn * lat_min_column, + const IColumn * lon_max_column, + const IColumn * lat_max_column, + const IColumn * precision_column, + ColumnPtr & result, + size_t input_rows_count) const { static constexpr size_t max_array_size = 10'000'000; + const auto * lon_min_const = typeid_cast(lon_min_column); + const auto * lat_min_const = typeid_cast(lat_min_column); + const auto * lon_max_const = typeid_cast(lon_max_column); + const auto * lat_max_const = typeid_cast(lat_max_column); + const auto * precision_const = typeid_cast(precision_column); + + if (lon_min_const) + lon_min_column = &lon_min_const->getDataColumn(); + if (lat_min_const) + lat_min_column = &lat_min_const->getDataColumn(); + if (lon_max_const) + lon_max_column = &lon_max_const->getDataColumn(); + if (lat_max_const) + lat_max_column = &lat_max_const->getDataColumn(); + if (precision_const) + precision_column = &precision_const->getDataColumn(); + const auto * lon_min = checkAndGetColumn>(lon_min_column); const auto * lat_min = checkAndGetColumn>(lat_min_column); const auto * lon_max = checkAndGetColumn>(lon_max_column); const auto * lat_max = checkAndGetColumn>(lat_max_column); - auto * precision = checkAndGetColumn>(precision_column); - if (precision == nullptr) - { - precision = checkAndGetColumnConstData>(precision_column); - } + const auto * precision = checkAndGetColumn>(precision_column); if (!lon_min || !lat_min || !lon_max || !lat_max || !precision) { @@ -88,24 +103,24 @@ public: ErrorCodes::LOGICAL_ERROR); } - const size_t total_rows = lat_min->size(); - auto col_res = ColumnArray::create(ColumnString::create()); ColumnString & res_strings = typeid_cast(col_res->getData()); ColumnArray::Offsets & res_offsets = col_res->getOffsets(); ColumnString::Chars & res_strings_chars = res_strings.getChars(); ColumnString::Offsets & res_strings_offsets = res_strings.getOffsets(); - for (size_t row = 0; row < total_rows; ++row) + for (size_t row = 0; row < input_rows_count; ++row) { - const Float64 lon_min_value = lon_min->getElement(row); - const Float64 lat_min_value = lat_min->getElement(row); - const Float64 lon_max_value = lon_max->getElement(row); - const Float64 lat_max_value = lat_max->getElement(row); + const Float64 lon_min_value = lon_min->getElement(lon_min_const ? 0 : row); + const Float64 lat_min_value = lat_min->getElement(lat_min_const ? 0 : row); + const Float64 lon_max_value = lon_max->getElement(lon_max_const ? 0 : row); + const Float64 lat_max_value = lat_max->getElement(lat_max_const ? 0 : row); + const PrecisionType precision_value = precision->getElement(precision_const ? 0 : row); const auto prepared_args = geohashesInBoxPrepare( - lon_min_value, lat_min_value, lon_max_value, lat_max_value, - precision->getElement(row % precision->size())); + lon_min_value, lat_min_value, lon_max_value, lat_max_value, + precision_value); + if (prepared_args.items_count > max_array_size) { throw Exception(getName() + " would produce " + std::to_string(prepared_args.items_count) + @@ -123,8 +138,9 @@ public: for (UInt64 i = 1; i <= prepared_args.items_count ; ++i) res_strings_offsets.push_back(starting_offset + (prepared_args.precision + 1) * i); - res_offsets.push_back((res_offsets.empty() ? 0 : res_offsets.back()) + prepared_args.items_count); + res_offsets.push_back(res_offsets.back() + prepared_args.items_count); } + if (!res_strings_offsets.empty() && res_strings_offsets.back() != res_strings_chars.size()) { throw Exception("String column size mismatch (internal logical error)", ErrorCodes::LOGICAL_ERROR); @@ -140,23 +156,19 @@ public: result = std::move(col_res); } - void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t /*input_rows_count*/) const override + void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) const override { const IColumn * lon_min = block.getByPosition(arguments[0]).column.get(); const IColumn * lat_min = block.getByPosition(arguments[1]).column.get(); const IColumn * lon_max = block.getByPosition(arguments[2]).column.get(); const IColumn * lat_max = block.getByPosition(arguments[3]).column.get(); - const IColumn * prec = block.getByPosition(arguments[4]).column.get(); + const IColumn * precision = block.getByPosition(arguments[4]).column.get(); ColumnPtr & res = block.getByPosition(result).column; if (checkColumn>(lon_min)) - { - execute(lon_min, lat_min, lon_max, lat_max, prec, res); - } + execute(lon_min, lat_min, lon_max, lat_max, precision, res, input_rows_count); else - { - execute(lon_min, lat_min, lon_max, lat_max, prec, res); - } + execute(lon_min, lat_min, lon_max, lat_max, precision, res, input_rows_count); } }; diff --git a/tests/queries/0_stateless/01426_geohash_constants.reference b/tests/queries/0_stateless/01426_geohash_constants.reference new file mode 100644 index 00000000000..ac62d60ce04 --- /dev/null +++ b/tests/queries/0_stateless/01426_geohash_constants.reference @@ -0,0 +1,6 @@ +['s'] +['s0'] +['s02','s08','s03','s09','s06','s0d'] +['s'] +['s0'] +['s02','s08','s03','s09','s06','s0d'] diff --git a/tests/queries/0_stateless/01426_geohash_constants.sql b/tests/queries/0_stateless/01426_geohash_constants.sql new file mode 100644 index 00000000000..4836ed6b2f8 --- /dev/null +++ b/tests/queries/0_stateless/01426_geohash_constants.sql @@ -0,0 +1,6 @@ +SELECT geohashesInBox(1., 2., 3., 4., 1); +SELECT geohashesInBox(materialize(1.), 2., 3., 4., 2); +SELECT geohashesInBox(1., materialize(2.), 3., 4., 3); +SELECT geohashesInBox(1., 2., materialize(3.), 4., 1); +SELECT geohashesInBox(1., 2., 3., materialize(4.), 2); +SELECT geohashesInBox(1., 2., 3., 4., materialize(3));