Separate star from ranges for better code

src/Functions/JSONPath/ASTs/ASTJSONPathStar.h

@@ -0,0 +1,15 @@
+#pragma once
+
+#include <Parsers/IAST.h>
+
+namespace DB
+{
+class ASTJSONPathStar : public IAST
+{
+public:
+    String getID(char) const override { return "ASTJSONPathStar"; }
+
+    ASTPtr clone() const override { return std::make_shared<ASTJSONPathStar>(*this); }
+};
+
+}

src/Functions/JSONPath/Generators/GeneratorJSONPath.h

@@ -4,6 +4,7 @@
 #include <Functions/JSONPath/Generators/VisitorJSONPathMemberAccess.h>
 #include <Functions/JSONPath/Generators/VisitorJSONPathRange.h>
 #include <Functions/JSONPath/Generators/VisitorJSONPathRoot.h>
+#include <Functions/JSONPath/Generators/VisitorJSONPathStar.h>
 #include <Functions/JSONPath/Generators/VisitorStatus.h>
 
 #include <Functions/JSONPath/ASTs/ASTJSONPath.h>
@@ -36,7 +37,8 @@ public:
 
         for (auto child_ast : query->children)
         {
-            if (typeid_cast<ASTJSONPathRoot *>(child_ast.get())) {
+            if (typeid_cast<ASTJSONPathRoot *>(child_ast.get()))
+            {
                 visitors.push_back(std::make_shared<VisitorJSONPathRoot<JSONParser>>(child_ast));
             }
             else if (typeid_cast<ASTJSONPathMemberAccess *>(child_ast.get()))
@@ -47,6 +49,10 @@ public:
             {
                 visitors.push_back(std::make_shared<VisitorJSONPathRange<JSONParser>>(child_ast));
             }
+            else if (typeid_cast<ASTJSONPathStar *>(child_ast.get()))
+            {
+                visitors.push_back(std::make_shared<VisitorJSONPathStar<JSONParser>>(child_ast));
+            }
         }
     }
 

src/Functions/JSONPath/Generators/VisitorJSONPathMemberAccess.h

@@ -10,15 +10,15 @@ template <typename JSONParser>
 class VisitorJSONPathMemberAccess : public IVisitor<JSONParser>
 {
 public:
-    VisitorJSONPathMemberAccess(ASTPtr member_access_ptr_) : member_access_ptr(member_access_ptr_) { }
+    VisitorJSONPathMemberAccess(ASTPtr member_access_ptr_)
+        : member_access_ptr(member_access_ptr_->as<ASTJSONPathMemberAccess>()) { }
 
     const char * getName() const override { return "VisitorJSONPathMemberAccess"; }
 
     VisitorStatus apply(typename JSONParser::Element & element) const override
     {
-        const auto * member_access = member_access_ptr->as<ASTJSONPathMemberAccess>();
         typename JSONParser::Element result;
-        element.getObject().find(std::string_view(member_access->member_name), result);
+        element.getObject().find(std::string_view(member_access_ptr->member_name), result);
         element = result;
         return VisitorStatus::Ok;
     }
@@ -30,9 +30,8 @@ public:
             this->setExhausted(true);
             return VisitorStatus::Error;
         }
-        const auto * member_access = member_access_ptr->as<ASTJSONPathMemberAccess>();
         typename JSONParser::Element result;
-        if (!element.getObject().find(std::string_view(member_access->member_name), result))
+        if (!element.getObject().find(std::string_view(member_access_ptr->member_name), result))
         {
             this->setExhausted(true);
             return VisitorStatus::Error;
@@ -47,7 +46,7 @@ public:
     void updateState() override { }
 
 private:
-    ASTPtr member_access_ptr;
+    ASTJSONPathMemberAccess * member_access_ptr;
 };
 
 }

src/Functions/JSONPath/Generators/VisitorJSONPathRange.h

@@ -10,18 +10,10 @@ template <typename JSONParser>
 class VisitorJSONPathRange : public IVisitor<JSONParser>
 {
 public:
-    VisitorJSONPathRange(ASTPtr range_ptr_) : range_ptr(range_ptr_)
+    VisitorJSONPathRange(ASTPtr range_ptr_) : range_ptr(range_ptr_->as<ASTJSONPathRange>())
     {
-        const auto * range = range_ptr->as<ASTJSONPathRange>();
         current_range = 0;
-        if (range->is_star)
-        {
-            current_index = 0;
-        }
-        else
-        {
-            current_index = range->ranges[current_range].first;
-        }
+        current_index = range_ptr->ranges[current_range].first;
     }
 
     const char * getName() const override { return "VisitorJSONPathRange"; }
@@ -30,12 +22,7 @@ public:
     {
         typename JSONParser::Element result;
         typename JSONParser::Array array = element.getArray();
-        if (current_index >= array.size())
-        {
-            return VisitorStatus::Error;
-        }
-        result = array[current_index];
-        element = result;
+        element = array[current_index];
         return VisitorStatus::Ok;
     }
 
@@ -47,69 +34,45 @@ public:
             return VisitorStatus::Error;
         }
 
-        const auto * range = range_ptr->as<ASTJSONPathRange>();
         VisitorStatus status;
         if (current_index < element.getArray().size())
         {
             apply(element);
             status = VisitorStatus::Ok;
         }
-        else if (!range->is_star)
-        {
-            status = VisitorStatus::Ignore;
-        }
         else
         {
             status = VisitorStatus::Ignore;
-            this->setExhausted(true);
         }
 
-        if (!range->is_star)
-        {
-            if (current_index + 1 == range->ranges[current_range].second)
-            {
-                if (current_range + 1 == range->ranges.size())
+        if (current_index + 1 == range_ptr->ranges[current_range].second
+            && current_range + 1 == range_ptr->ranges.size())
         {
             this->setExhausted(true);
         }
-            }
-        }
 
         return status;
     }
 
     void reinitialize() override
     {
-        const auto * range = range_ptr->as<ASTJSONPathRange>();
         current_range = 0;
-        if (range->is_star)
-        {
-            current_index = 0;
-        }
-        else
-        {
-            current_index = range->ranges[current_range].first;
-        }
+        current_index = range_ptr->ranges[current_range].first;
         this->setExhausted(false);
     }
 
     void updateState() override
     {
-        const auto * range = range_ptr->as<ASTJSONPathRange>();
         current_index++;
-        if (range->is_star)
-        {
-            return;
-        }
-        if (current_index == range->ranges[current_range].second)
+        if (current_index == range_ptr->ranges[current_range].second)
         {
             current_range++;
-            current_index = range->ranges[current_range].first;
+            current_index = range_ptr->ranges[current_range].first;
         }
     }
 
 private:
-    ASTPtr range_ptr;
+    ASTJSONPathRange * range_ptr;
     size_t current_range;
     UInt32 current_index;
 };

src/Functions/JSONPath/Generators/VisitorJSONPathStar.h

@@ -0,0 +1,66 @@
+#pragma once
+
+#include <Functions/JSONPath/ASTs/ASTJSONPathStar.h>
+#include <Functions/JSONPath/Generators/IVisitor.h>
+#include <Functions/JSONPath/Generators/VisitorStatus.h>
+
+namespace DB
+{
+template <typename JSONParser>
+class VisitorJSONPathStar : public IVisitor<JSONParser>
+{
+public:
+    VisitorJSONPathStar(ASTPtr)
+    {
+        current_index = 0;
+    }
+
+    const char * getName() const override { return "VisitorJSONPathStar"; }
+
+    VisitorStatus apply(typename JSONParser::Element & element) const override
+    {
+        typename JSONParser::Element result;
+        typename JSONParser::Array array = element.getArray();
+        element = array[current_index];
+        return VisitorStatus::Ok;
+    }
+
+    VisitorStatus visit(typename JSONParser::Element & element) override
+    {
+        if (!element.isArray())
+        {
+            this->setExhausted(true);
+            return VisitorStatus::Error;
+        }
+
+        VisitorStatus status;
+        if (current_index < element.getArray().size())
+        {
+            apply(element);
+            status = VisitorStatus::Ok;
+        }
+        else
+        {
+            status = VisitorStatus::Ignore;
+            this->setExhausted(true);
+        }
+
+        return status;
+    }
+
+    void reinitialize() override
+    {
+        current_index = 0;
+        this->setExhausted(false);
+    }
+
+    void updateState() override
+    {
+        current_index++;
+    }
+
+private:
+    UInt32 current_index;
+};
+
+}

src/Functions/JSONPath/Parsers/ParserJSONPathQuery.cpp

@@ -3,6 +3,7 @@
 #include <Functions/JSONPath/Parsers/ParserJSONPathRoot.h>
 #include <Functions/JSONPath/Parsers/ParserJSONPathMemberAccess.h>
 #include <Functions/JSONPath/Parsers/ParserJSONPathRange.h>
+#include <Functions/JSONPath/Parsers/ParserJSONPathStar.h>
 
 namespace DB
 
@@ -19,17 +20,20 @@ bool ParserJSONPathQuery::parseImpl(Pos & pos, ASTPtr & query, Expected & expect
     query = std::make_shared<ASTJSONPathQuery>();
     ParserJSONPathMemberAccess parser_jsonpath_member_access;
     ParserJSONPathRange parser_jsonpath_range;
+    ParserJSONPathStar parser_jsonpath_star;
     ParserJSONPathRoot parser_jsonpath_root;
 
     ASTPtr path_root;
-    if (!parser_jsonpath_root.parse(pos, path_root, expected)) {
+    if (!parser_jsonpath_root.parse(pos, path_root, expected))
+    {
         return false;
     }
     query->children.push_back(path_root);
 
     ASTPtr accessor;
     while (parser_jsonpath_member_access.parse(pos, accessor, expected)
-           || parser_jsonpath_range.parse(pos, accessor, expected))
+           || parser_jsonpath_range.parse(pos, accessor, expected)
+           || parser_jsonpath_star.parse(pos, accessor, expected))
     {
         if (accessor)
         {

src/Functions/JSONPath/Parsers/ParserJSONPathRange.cpp

@@ -5,6 +5,7 @@
 #include <Parsers/ASTIdentifier.h>
 #include <Parsers/ASTLiteral.h>
 #include <Parsers/ExpressionElementParsers.h>
+#include <Parsers/CommonParsers.h>
 
 namespace DB
 {
@@ -31,26 +32,16 @@ bool ParserJSONPathRange::parseImpl(Pos & pos, ASTPtr & node, Expected & expecte
     auto range = std::make_shared<ASTJSONPathRange>();
     node = range;
 
+    ParserNumber number_p;
+    ASTPtr number_ptr;
     while (pos->type != TokenType::ClosingSquareBracket)
     {
-        if (pos->type != TokenType::Number && pos->type != TokenType::Asterisk)
+        if (pos->type != TokenType::Number)
         {
             return false;
         }
-        if (pos->type == TokenType::Asterisk)
-        {
-            if (range->is_star)
-            {
-                throw Exception("Multiple asterisks in square array range are not allowed", ErrorCodes::BAD_ARGUMENTS);
-            }
-            range->is_star = true;
-            ++pos;
-            continue;
-        }
 
         std::pair<UInt32, UInt32> range_indices;
-        ParserNumber number_p;
-        ASTPtr number_ptr;
         if (!number_p.parse(pos, number_ptr, expected))
         {
             return false;
@@ -64,16 +55,7 @@ bool ParserJSONPathRange::parseImpl(Pos & pos, ASTPtr & node, Expected & expecte
         }
         else if (pos->type == TokenType::BareWord)
         {
-            /// Range case
-            ParserIdentifier name_p;
-            ASTPtr word;
-            if (!name_p.parse(pos, word, expected))
-            {
-                return false;
-            }
-            String to_identifier;
-            if (!tryGetIdentifierNameInto(word, to_identifier) || to_identifier != "to")
-            {
+            if (!ParserKeyword("TO").ignore(pos, expected)) {
                 return false;
             }
             if (!number_p.parse(pos, number_ptr, expected))

src/Functions/JSONPath/Parsers/ParserJSONPathStar.cpp

@@ -0,0 +1,31 @@
+#include <Functions/JSONPath/Parsers/ParserJSONPathStar.h>
+
+#include <Functions/JSONPath/ASTs/ASTJSONPathStar.h>
+
+namespace DB
+{
+bool ParserJSONPathStar::parseImpl(Pos & pos, ASTPtr & node, Expected & expected)
+{
+
+    if (pos->type != TokenType::OpeningSquareBracket)
+    {
+        return false;
+    }
+    ++pos;
+    if (pos->type != TokenType::Asterisk) {
+        return false;
+    }
+    ++pos;
+    if (pos->type != TokenType::ClosingSquareBracket) {
+        expected.add(pos, "Closing square bracket");
+        return false;
+    }
+    ++pos;
+
+    auto star = std::make_shared<ASTJSONPathStar>();
+    node = star;
+
+    return true;
+}
+
+}

src/Functions/JSONPath/Parsers/ParserJSONPathStar.h

@@ -0,0 +1,18 @@
+#pragma once
+
+#include <Parsers/IParserBase.h>
+
+
+namespace DB
+{
+class ParserJSONPathStar : public IParserBase
+{
+private:
+    const char * getName() const override { return "ParserJSONPathStar"; }
+    bool parseImpl(Pos & pos, ASTPtr & node, Expected & expected) override;
+
+public:
+    explicit ParserJSONPathStar() = default;
+};
+
+}