jsonpath: normalize unary operations on numeric scalars

This commit adds normalization for unary plus and minus operations on
numeric scalar values in JSONPath expressions during parsing. This
increases Postgres compatibility.

Release note: None

Author: normanchenn

pkg/sql/scanner/jsonpath_scan.go

@@ -18,7 +18,7 @@ type JSONPathScanner struct {
 // Scan scans the next token and populates its information into lval.
 // This scan function contains rules for jsonpath.
 func (s *JSONPathScanner) Scan(lval ScanSymType) {
-	ch, skipWhiteSpace := s.scanSetup(lval)
+	ch, skipWhiteSpace := s.scanSetup(lval, false /* allowComments */)
 	if skipWhiteSpace {
 		return
 	}

pkg/sql/scanner/plpgsql_scan.go

@@ -21,7 +21,7 @@ type PLpgSQLScanner struct {
 // Scan scans the next token and populates its information into lval.
 // This scan function contains rules for plpgsql.
 func (s *PLpgSQLScanner) Scan(lval ScanSymType) {
-	ch, skipWhiteSpace := s.scanSetup(lval)
+	ch, skipWhiteSpace := s.scanSetup(lval, true /* allowComments */)
 
 	if skipWhiteSpace {
 		return

pkg/sql/scanner/scan.go

@@ -141,14 +141,14 @@ func (s *Scanner) finishString(buf []byte) string {
 	return str
 }
 
-func (s *Scanner) scanSetup(lval ScanSymType) (int, bool) {
+func (s *Scanner) scanSetup(lval ScanSymType, allowComments bool) (int, bool) {
 	lval.SetID(0)
 	lval.SetPos(int32(s.pos))
 	lval.SetStr("EOF")
 	s.quoted = false
 	s.lastAttemptedID = 0
 
-	if _, ok := s.skipWhitespace(lval, true); !ok {
+	if _, ok := s.skipWhitespace(lval, allowComments); !ok {
 		return 0, true
 	}
 
@@ -167,7 +167,7 @@ func (s *Scanner) scanSetup(lval ScanSymType) (int, bool) {
 
 // Scan scans the next token and populates its information into lval.
 func (s *SQLScanner) Scan(lval ScanSymType) {
-	ch, skipWhiteSpace := s.scanSetup(lval)
+	ch, skipWhiteSpace := s.scanSetup(lval, true /* allowComments */)
 
 	if skipWhiteSpace {
 		return

pkg/util/jsonpath/operation.go

@@ -70,13 +70,6 @@ func (o Operation) String() string {
 	if o.Type == OpLogicalNot {
 		return fmt.Sprintf("%s(%s)", OperationTypeStrings[o.Type], o.Left)
 	}
-	// TODO(normanchenn): Postgres normalizes unary +/- operators differently
-	// for numbers vs. non-numbers.
-	// Numbers:      '-1' -> '-1', '--1' -> '1'
-	// Non-numbers:  '-"hello"' -> '(-"hello")'
-	// We currently don't normalize numbers - we output `(-1)` and `(-(-1))`.
-	// See makeItemUnary in postgres/src/backend/utils/adt/jsonpath_gram.y. This
-	// can be done at parse time.
 	if o.Type == OpPlus || o.Type == OpMinus {
 		return fmt.Sprintf("(%s%s)", OperationTypeStrings[o.Type], o.Left)
 	}

pkg/util/jsonpath/parser/jsonpath.y

@@ -119,14 +119,59 @@ func binaryOp(op jsonpath.OperationType, left jsonpath.Path, right jsonpath.Path
   }
 }
 
-func unaryOp(op jsonpath.OperationType, left jsonpath.Path) jsonpath.Operation {
+func unaryOp(op jsonpath.OperationType, left jsonpath.Path) jsonpath.Path {
+  if scalar, ok := maybeNormalizeUnaryOp(op, left); ok {
+    return scalar
+  }
   return jsonpath.Operation{
     Type:  op,
     Left:  left,
     Right: nil,
   }
 }
 
+func maybeNormalizeUnaryOp(op jsonpath.OperationType, expr jsonpath.Path) (jsonpath.Scalar, bool) {
+  // Return if not unary plus or unary minus.
+  if op != jsonpath.OpPlus && op != jsonpath.OpMinus {
+    return jsonpath.Scalar{}, false
+  }
+
+  scalar, ok := extractNumericScalar(expr)
+  if !ok {
+    return jsonpath.Scalar{}, false
+  }
+  if op == jsonpath.OpMinus {
+    dec, _ := scalar.Value.AsDecimal()
+    dec.Neg(dec)
+    scalar.Value = json.FromDecimal(*dec)
+  }
+  return scalar, true
+}
+
+// extractNumericScalar attempts to extract a numeric scalar value from a
+// jsonpath.Path. It handles two cases:
+// - Direct scalar values.
+// - Scalar values wrapped in a jsonpath.Paths object of length 1.
+// The function returns the scalar and true if the path contains a valid numeric
+// value (integer or float), otherwise returns an empty scalar and false.
+func extractNumericScalar(expr jsonpath.Path) (jsonpath.Scalar, bool) {
+  potentialScalar := expr
+  if paths, ok := expr.(jsonpath.Paths); ok {
+    if len(paths) != 1 {
+      return jsonpath.Scalar{}, false
+    }
+    potentialScalar = paths[0]
+  }
+  scalar, ok := potentialScalar.(jsonpath.Scalar)
+  if !ok {
+    return jsonpath.Scalar{}, false
+  }
+  if scalar.Type != jsonpath.ScalarFloat && scalar.Type != jsonpath.ScalarInt {
+    return jsonpath.Scalar{}, false
+  }
+  return scalar, true
+}
+
 func regexBinaryOp(left jsonpath.Path, regex string) (jsonpath.Operation, error) {
   r := jsonpath.Regex{Regex: regex}
   _, err := ReCache.GetRegexp(r)

pkg/util/jsonpath/parser/testdata/jsonpath

@@ -471,7 +471,7 @@ $."abc".*."def".* -- normalized!
 parse
 -1
 ----
-(-1) -- normalized!
+-1
 
 parse
 - "hello"
@@ -491,7 +491,7 @@ parse
 parse
 1 + 2 * -4
 ----
-(1 + (2 * (-4))) -- normalized!
+(1 + (2 * -4)) -- normalized!
 
 error
 $ ? (@ like_regex "(invalid pattern")
@@ -622,16 +622,15 @@ $.abc[0.0]
 ----
 $."abc"[0] -- normalized!
 
-# TODO(normanchenn): Related to the TODOs in pkg/util/jsonpath/operation.go.
 parse
 $.abc[-0]
 ----
-$."abc"[(-0)] -- normalized!
+$."abc"[0] -- normalized!
 
 parse
 $.abc[-1.99]
 ----
-$."abc"[(-1.99)] -- normalized!
+$."abc"[-1.99] -- normalized!
 
 parse
 $[1.999999999999999]
@@ -643,6 +642,46 @@ $[null]
 ----
 $[null]
 
+parse
++1
+----
+1 -- normalized!
+
+parse
+++1
+----
+1 -- normalized!
+
+parse
++++1
+----
+1 -- normalized!
+
+parse
+- -1
+----
+1 -- normalized!
+
+parse
+- - -1
+----
+-1 -- normalized!
+
+parse
+(1--123).type()
+----
+(1 - -123).type() -- normalized!
+
+parse
+(1-+-123).type()
+----
+(1 - -123).type() -- normalized!
+
+parse
+(1+-+-123).type()
+----
+(1 + 123).type() -- normalized!
+
 # parse
 # $.1a
 # ----