github.com/whtcorpsinc/milevadb-prod@v0.0.0-20211104133533-f57f4be3b597/dbs/memristed/memex/builtin_json_vec.go (about) 1 // Copyright 2020 WHTCORPS INC, Inc. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // See the License for the specific language governing permissions and 12 // limitations under the License. 13 14 package memex 15 16 import ( 17 "strconv" 18 "strings" 19 20 "github.com/whtcorpsinc/errors" 21 "github.com/whtcorpsinc/BerolinaSQL/ast" 22 "github.com/whtcorpsinc/milevadb/stochastikctx" 23 "github.com/whtcorpsinc/milevadb/types" 24 "github.com/whtcorpsinc/milevadb/types/json" 25 "github.com/whtcorpsinc/milevadb/soliton/chunk" 26 "github.com/whtcorpsinc/fidelpb/go-fidelpb" 27 ) 28 29 func vecJSONModify(ctx stochastikctx.Context, args []Expression, bufSlabPredictor defCausumnBufferSlabPredictor, input *chunk.Chunk, result *chunk.DeferredCauset, mt json.ModifyType) error { 30 nr := input.NumEvents() 31 jsonBuf, err := bufSlabPredictor.get(types.ETJson, nr) 32 if err != nil { 33 return err 34 } 35 defer bufSlabPredictor.put(jsonBuf) 36 if err := args[0].VecEvalJSON(ctx, input, jsonBuf); err != nil { 37 return err 38 } 39 40 strBufs := make([]*chunk.DeferredCauset, (len(args)-1)/2) 41 for i := 1; i < len(args); i += 2 { 42 strBufs[(i-1)/2], err = bufSlabPredictor.get(types.ETString, nr) 43 if err != nil { 44 return err 45 } 46 defer bufSlabPredictor.put(strBufs[(i-1)/2]) 47 if err := args[i].VecEvalString(ctx, input, strBufs[(i-1)/2]); err != nil { 48 return err 49 } 50 } 51 valueBufs := make([]*chunk.DeferredCauset, (len(args)-1)/2+1) 52 for i := 2; i < len(args); i += 2 { 53 valueBufs[i/2-1], err = bufSlabPredictor.get(types.ETJson, nr) 54 if err != nil { 55 return err 56 } 57 defer bufSlabPredictor.put(valueBufs[i/2-1]) 58 if err := args[i].VecEvalJSON(ctx, input, valueBufs[i/2-1]); err != nil { 59 return err 60 } 61 } 62 result.ReserveJSON(nr) 63 for i := 0; i < nr; i++ { 64 if jsonBuf.IsNull(i) { 65 result.AppendNull() 66 continue 67 } 68 pathExprs := make([]json.PathExpression, 0, (len(args)-1)/2+1) 69 values := make([]json.BinaryJSON, 0, (len(args)-1)/2+1) 70 var pathExpr json.PathExpression 71 isNull := false 72 for j := 1; j < len(args); j += 2 { 73 if strBufs[(j-1)/2].IsNull(i) { 74 isNull = true 75 break 76 } 77 pathExpr, err = json.ParseJSONPathExpr(strBufs[(j-1)/2].GetString(i)) 78 if err != nil { 79 return err 80 } 81 pathExprs = append(pathExprs, pathExpr) 82 } 83 for j := 2; j < len(args); j += 2 { 84 if valueBufs[j/2-1].IsNull(i) { 85 values = append(values, json.CreateBinary(nil)) 86 } else { 87 values = append(values, valueBufs[j/2-1].GetJSON(i)) 88 } 89 } 90 if isNull { 91 result.AppendNull() 92 } else { 93 res, err := jsonBuf.GetJSON(i).Modify(pathExprs, values, mt) 94 if err != nil { 95 return err 96 } 97 result.AppendJSON(res) 98 } 99 } 100 return nil 101 } 102 103 func (b *builtinJSONStorageSizeSig) vectorized() bool { 104 return true 105 } 106 107 func (b *builtinJSONStorageSizeSig) vecEvalInt(input *chunk.Chunk, result *chunk.DeferredCauset) error { 108 n := input.NumEvents() 109 buf, err := b.bufSlabPredictor.get(types.ETJson, n) 110 if err != nil { 111 return err 112 } 113 defer b.bufSlabPredictor.put(buf) 114 if err := b.args[0].VecEvalJSON(b.ctx, input, buf); err != nil { 115 return err 116 } 117 result.ResizeInt64(n, false) 118 result.MergeNulls(buf) 119 int64s := result.Int64s() 120 for i := 0; i < n; i++ { 121 if result.IsNull(i) { 122 continue 123 } 124 j := buf.GetJSON(i) 125 126 jb, err := j.MarshalJSON() 127 if err != nil { 128 continue 129 } 130 131 int64s[i] = int64(len(jb)) 132 } 133 return nil 134 } 135 136 func (b *builtinJSONDepthSig) vectorized() bool { 137 return true 138 } 139 140 func (b *builtinJSONDepthSig) vecEvalInt(input *chunk.Chunk, result *chunk.DeferredCauset) error { 141 n := input.NumEvents() 142 buf, err := b.bufSlabPredictor.get(types.ETJson, n) 143 if err != nil { 144 return err 145 } 146 defer b.bufSlabPredictor.put(buf) 147 if err := b.args[0].VecEvalJSON(b.ctx, input, buf); err != nil { 148 return err 149 } 150 result.ResizeInt64(n, false) 151 result.MergeNulls(buf) 152 int64s := result.Int64s() 153 for i := 0; i < n; i++ { 154 if result.IsNull(i) { 155 continue 156 } 157 j := buf.GetJSON(i) 158 int64s[i] = int64(j.GetElemDepth()) 159 } 160 return nil 161 } 162 163 func (b *builtinJSONKeysSig) vectorized() bool { 164 return true 165 } 166 167 func (b *builtinJSONKeysSig) vecEvalJSON(input *chunk.Chunk, result *chunk.DeferredCauset) error { 168 n := input.NumEvents() 169 buf, err := b.bufSlabPredictor.get(types.ETJson, n) 170 if err != nil { 171 return err 172 } 173 defer b.bufSlabPredictor.put(buf) 174 if err := b.args[0].VecEvalJSON(b.ctx, input, buf); err != nil { 175 return err 176 } 177 178 result.ReserveJSON(n) 179 var j json.BinaryJSON 180 for i := 0; i < n; i++ { 181 if buf.IsNull(i) { 182 result.AppendNull() 183 continue 184 } 185 186 j = buf.GetJSON(i) 187 if j.TypeCode != json.TypeCodeObject { 188 result.AppendNull() 189 continue 190 } 191 result.AppendJSON(j.GetKeys()) 192 } 193 return nil 194 } 195 196 func (b *builtinJSONInsertSig) vectorized() bool { 197 return true 198 } 199 200 func (b *builtinJSONInsertSig) vecEvalJSON(input *chunk.Chunk, result *chunk.DeferredCauset) error { 201 err := vecJSONModify(b.ctx, b.args, b.bufSlabPredictor, input, result, json.ModifyInsert) 202 return err 203 } 204 205 func (b *builtinJSONReplaceSig) vectorized() bool { 206 return true 207 } 208 209 func (b *builtinJSONReplaceSig) vecEvalJSON(input *chunk.Chunk, result *chunk.DeferredCauset) error { 210 err := vecJSONModify(b.ctx, b.args, b.bufSlabPredictor, input, result, json.ModifyReplace) 211 return err 212 } 213 214 func (b *builtinJSONArraySig) vectorized() bool { 215 return true 216 } 217 218 func (b *builtinJSONArraySig) vecEvalJSON(input *chunk.Chunk, result *chunk.DeferredCauset) error { 219 nr := input.NumEvents() 220 jsons := make([][]interface{}, nr) 221 for i := 0; i < nr; i++ { 222 jsons[i] = make([]interface{}, 0, len(b.args)) 223 } 224 for _, arg := range b.args { 225 j, err := b.bufSlabPredictor.get(types.ETJson, nr) 226 if err != nil { 227 return err 228 } 229 defer b.bufSlabPredictor.put(j) 230 if err = arg.VecEvalJSON(b.ctx, input, j); err != nil { 231 return err 232 } 233 for i := 0; i < nr; i++ { 234 if j.IsNull(i) { 235 jsons[i] = append(jsons[i], json.CreateBinary(nil)) 236 } else { 237 jsons[i] = append(jsons[i], j.GetJSON(i)) 238 } 239 } 240 } 241 result.ReserveJSON(nr) 242 for i := 0; i < nr; i++ { 243 result.AppendJSON(json.CreateBinary(jsons[i])) 244 } 245 return nil 246 } 247 248 func (b *builtinJSONContainsSig) vectorized() bool { 249 return true 250 } 251 252 func (b *builtinJSONContainsSig) vecEvalInt(input *chunk.Chunk, result *chunk.DeferredCauset) error { 253 nr := input.NumEvents() 254 255 objDefCaus, err := b.bufSlabPredictor.get(types.ETJson, nr) 256 if err != nil { 257 return err 258 } 259 defer b.bufSlabPredictor.put(objDefCaus) 260 261 if err := b.args[0].VecEvalJSON(b.ctx, input, objDefCaus); err != nil { 262 return err 263 } 264 265 targetDefCaus, err := b.bufSlabPredictor.get(types.ETJson, nr) 266 if err != nil { 267 return err 268 } 269 defer b.bufSlabPredictor.put(targetDefCaus) 270 271 if err := b.args[1].VecEvalJSON(b.ctx, input, targetDefCaus); err != nil { 272 return err 273 } 274 275 result.ResizeInt64(nr, false) 276 resI64s := result.Int64s() 277 278 if len(b.args) == 3 { 279 pathDefCaus, err := b.bufSlabPredictor.get(types.ETString, nr) 280 if err != nil { 281 return err 282 } 283 defer b.bufSlabPredictor.put(pathDefCaus) 284 285 if err := b.args[2].VecEvalString(b.ctx, input, pathDefCaus); err != nil { 286 return err 287 } 288 289 result.MergeNulls(objDefCaus, targetDefCaus, pathDefCaus) 290 291 var pathExpr json.PathExpression 292 for i := 0; i < nr; i++ { 293 if result.IsNull(i) { 294 continue 295 } 296 pathExpr, err = json.ParseJSONPathExpr(pathDefCaus.GetString(i)) 297 if err != nil { 298 return err 299 } 300 if pathExpr.ContainsAnyAsterisk() { 301 return json.ErrInvalidJSONPathWildcard 302 } 303 304 obj, exists := objDefCaus.GetJSON(i).Extract([]json.PathExpression{pathExpr}) 305 if !exists { 306 result.SetNull(i, true) 307 continue 308 } 309 310 if json.ContainsBinary(obj, targetDefCaus.GetJSON(i)) { 311 resI64s[i] = 1 312 } else { 313 resI64s[i] = 0 314 } 315 } 316 } else { 317 result.MergeNulls(objDefCaus, targetDefCaus) 318 for i := 0; i < nr; i++ { 319 if result.IsNull(i) { 320 continue 321 } 322 if json.ContainsBinary(objDefCaus.GetJSON(i), targetDefCaus.GetJSON(i)) { 323 resI64s[i] = 1 324 } else { 325 resI64s[i] = 0 326 } 327 } 328 } 329 330 return nil 331 } 332 333 func (b *builtinJSONQuoteSig) vectorized() bool { 334 return true 335 } 336 337 func (b *builtinJSONQuoteSig) vecEvalString(input *chunk.Chunk, result *chunk.DeferredCauset) error { 338 n := input.NumEvents() 339 buf, err := b.bufSlabPredictor.get(types.ETString, n) 340 if err != nil { 341 return err 342 } 343 defer b.bufSlabPredictor.put(buf) 344 if err := b.args[0].VecEvalString(b.ctx, input, buf); err != nil { 345 return err 346 } 347 348 result.ReserveString(n) 349 for i := 0; i < n; i++ { 350 if buf.IsNull(i) { 351 result.AppendNull() 352 continue 353 } 354 result.AppendString(strconv.Quote(buf.GetString(i))) 355 } 356 return nil 357 } 358 359 func (b *builtinJSONSearchSig) vectorized() bool { 360 return true 361 } 362 363 func (b *builtinJSONSearchSig) vecEvalJSON(input *chunk.Chunk, result *chunk.DeferredCauset) error { 364 nr := input.NumEvents() 365 jsonBuf, err := b.bufSlabPredictor.get(types.ETJson, nr) 366 if err != nil { 367 return err 368 } 369 defer b.bufSlabPredictor.put(jsonBuf) 370 if err := b.args[0].VecEvalJSON(b.ctx, input, jsonBuf); err != nil { 371 return err 372 } 373 typeBuf, err := b.bufSlabPredictor.get(types.ETString, nr) 374 if err != nil { 375 return err 376 } 377 defer b.bufSlabPredictor.put(typeBuf) 378 if err := b.args[1].VecEvalString(b.ctx, input, typeBuf); err != nil { 379 return err 380 } 381 searchBuf, err := b.bufSlabPredictor.get(types.ETString, nr) 382 if err != nil { 383 return err 384 } 385 defer b.bufSlabPredictor.put(searchBuf) 386 if err := b.args[2].VecEvalString(b.ctx, input, searchBuf); err != nil { 387 return err 388 } 389 390 var escapeBuf *chunk.DeferredCauset 391 if len(b.args) >= 4 { 392 escapeBuf, err = b.bufSlabPredictor.get(types.ETString, nr) 393 if err != nil { 394 return err 395 } 396 defer b.bufSlabPredictor.put(escapeBuf) 397 if err := b.args[3].VecEvalString(b.ctx, input, escapeBuf); err != nil { 398 return err 399 } 400 } 401 402 var pathBufs []*chunk.DeferredCauset 403 if len(b.args) >= 5 { 404 pathBufs = make([]*chunk.DeferredCauset, (len(b.args) - 4)) 405 for i := 4; i < len(b.args); i++ { 406 index := i - 4 407 pathBufs[index], err = b.bufSlabPredictor.get(types.ETString, nr) 408 if err != nil { 409 return err 410 } 411 defer b.bufSlabPredictor.put(pathBufs[index]) 412 if err := b.args[i].VecEvalString(b.ctx, input, pathBufs[index]); err != nil { 413 return err 414 } 415 } 416 } 417 418 result.ReserveJSON(nr) 419 420 for i := 0; i < nr; i++ { 421 if jsonBuf.IsNull(i) || searchBuf.IsNull(i) || typeBuf.IsNull(i) { 422 result.AppendNull() 423 continue 424 } 425 containType := strings.ToLower(typeBuf.GetString(i)) 426 escape := byte('\\') 427 if escapeBuf != nil && !escapeBuf.IsNull(i) { 428 escapeStr := escapeBuf.GetString(i) 429 if len(escapeStr) == 0 { 430 escape = byte('\\') 431 } else if len(escapeStr) == 1 { 432 escape = escapeStr[0] 433 } else { 434 return errIncorrectArgs.GenWithStackByArgs("ESCAPE") 435 } 436 } 437 var pathExprs []json.PathExpression 438 if pathBufs != nil { 439 pathExprs = make([]json.PathExpression, 0, len(b.args)-4) 440 for j := 0; j < len(b.args)-4; j++ { 441 if pathBufs[j].IsNull(i) { 442 break 443 } 444 pathExpr, err := json.ParseJSONPathExpr(pathBufs[j].GetString(i)) 445 if err != nil { 446 return json.ErrInvalidJSONPath.GenWithStackByArgs(pathBufs[j].GetString(i)) 447 } 448 pathExprs = append(pathExprs, pathExpr) 449 } 450 } 451 bj, _, err := jsonBuf.GetJSON(i).Search(containType, searchBuf.GetString(i), escape, pathExprs) 452 if err != nil { 453 return err 454 } 455 result.AppendJSON(bj) 456 } 457 return nil 458 } 459 460 func (b *builtinJSONSetSig) vectorized() bool { 461 return true 462 } 463 464 func (b *builtinJSONSetSig) vecEvalJSON(input *chunk.Chunk, result *chunk.DeferredCauset) error { 465 err := vecJSONModify(b.ctx, b.args, b.bufSlabPredictor, input, result, json.ModifySet) 466 return err 467 } 468 469 func (b *builtinJSONObjectSig) vectorized() bool { 470 return true 471 } 472 473 func (b *builtinJSONObjectSig) vecEvalJSON(input *chunk.Chunk, result *chunk.DeferredCauset) error { 474 nr := input.NumEvents() 475 if len(b.args)&1 == 1 { 476 err := ErrIncorrectParameterCount.GenWithStackByArgs(ast.JSONObject) 477 return err 478 } 479 480 jsons := make([]map[string]interface{}, nr) 481 for i := 0; i < nr; i++ { 482 jsons[i] = make(map[string]interface{}, len(b.args)>>1) 483 } 484 485 argBuffers := make([]*chunk.DeferredCauset, len(b.args)) 486 var err error 487 for i := 0; i < len(b.args); i++ { 488 if i&1 == 0 { 489 if argBuffers[i], err = b.bufSlabPredictor.get(types.ETString, nr); err != nil { 490 return err 491 } 492 defer func(buf *chunk.DeferredCauset) { 493 b.bufSlabPredictor.put(buf) 494 }(argBuffers[i]) 495 496 if err = b.args[i].VecEvalString(b.ctx, input, argBuffers[i]); err != nil { 497 return err 498 } 499 } else { 500 if argBuffers[i], err = b.bufSlabPredictor.get(types.ETJson, nr); err != nil { 501 return err 502 } 503 defer func(buf *chunk.DeferredCauset) { 504 b.bufSlabPredictor.put(buf) 505 }(argBuffers[i]) 506 507 if err = b.args[i].VecEvalJSON(b.ctx, input, argBuffers[i]); err != nil { 508 return err 509 } 510 } 511 } 512 513 result.ReserveJSON(nr) 514 for i := 0; i < len(b.args); i++ { 515 if i&1 == 1 { 516 keyDefCaus := argBuffers[i-1] 517 valueDefCaus := argBuffers[i] 518 519 var key string 520 var value json.BinaryJSON 521 for j := 0; j < nr; j++ { 522 if keyDefCaus.IsNull(j) { 523 err := errors.New("JSON documents may not contain NULL member names") 524 return err 525 } 526 key = keyDefCaus.GetString(j) 527 if valueDefCaus.IsNull(j) { 528 value = json.CreateBinary(nil) 529 } else { 530 value = valueDefCaus.GetJSON(j) 531 } 532 jsons[j][key] = value 533 } 534 } 535 } 536 537 for i := 0; i < nr; i++ { 538 result.AppendJSON(json.CreateBinary(jsons[i])) 539 } 540 return nil 541 } 542 543 func (b *builtinJSONArrayInsertSig) vectorized() bool { 544 return true 545 } 546 547 func (b *builtinJSONArrayInsertSig) vecEvalJSON(input *chunk.Chunk, result *chunk.DeferredCauset) error { 548 nr := input.NumEvents() 549 buf, err := b.bufSlabPredictor.get(types.ETJson, nr) 550 if err != nil { 551 return err 552 } 553 defer b.bufSlabPredictor.put(buf) 554 if err := b.args[0].VecEvalJSON(b.ctx, input, buf); err != nil { 555 return err 556 } 557 pathBufs := make([]*chunk.DeferredCauset, (len(b.args)-1)/2) 558 valueBufs := make([]*chunk.DeferredCauset, (len(b.args)-1)/2) 559 for i := 1; i < len(b.args); i++ { 560 if i&1 == 0 { 561 valueBufs[i/2-1], err = b.bufSlabPredictor.get(types.ETJson, nr) 562 if err != nil { 563 return err 564 } 565 defer b.bufSlabPredictor.put(valueBufs[i/2-1]) 566 if err := b.args[i].VecEvalJSON(b.ctx, input, valueBufs[i/2-1]); err != nil { 567 return err 568 } 569 } else { 570 pathBufs[(i-1)/2], err = b.bufSlabPredictor.get(types.ETString, nr) 571 if err != nil { 572 return err 573 } 574 defer b.bufSlabPredictor.put(pathBufs[(i-1)/2]) 575 if err := b.args[i].VecEvalString(b.ctx, input, pathBufs[(i-1)/2]); err != nil { 576 return err 577 } 578 } 579 } 580 var pathExpr json.PathExpression 581 var value json.BinaryJSON 582 result.ReserveJSON(nr) 583 for i := 0; i < nr; i++ { 584 if buf.IsNull(i) { 585 result.AppendNull() 586 continue 587 } 588 589 res := buf.GetJSON(i) 590 isnull := false 591 for j := 0; j < (len(b.args)-1)/2; j++ { 592 if pathBufs[j].IsNull(i) { 593 isnull = true 594 break 595 } 596 pathExpr, err = json.ParseJSONPathExpr(pathBufs[j].GetString(i)) 597 if err != nil { 598 return json.ErrInvalidJSONPath.GenWithStackByArgs(pathBufs[j].GetString(i)) 599 } 600 if pathExpr.ContainsAnyAsterisk() { 601 return json.ErrInvalidJSONPathWildcard.GenWithStackByArgs(pathBufs[j].GetString(i)) 602 } 603 if valueBufs[j].IsNull(i) { 604 value = json.CreateBinary(nil) 605 } else { 606 value = valueBufs[j].GetJSON(i) 607 } 608 res, err = res.ArrayInsert(pathExpr, value) 609 if err != nil { 610 return err 611 } 612 } 613 if isnull { 614 result.AppendNull() 615 continue 616 } 617 result.AppendJSON(res) 618 } 619 return nil 620 } 621 622 func (b *builtinJSONKeys2ArgsSig) vectorized() bool { 623 return true 624 } 625 626 func (b *builtinJSONKeys2ArgsSig) vecEvalJSON(input *chunk.Chunk, result *chunk.DeferredCauset) error { 627 nr := input.NumEvents() 628 jsonBuf, err := b.bufSlabPredictor.get(types.ETJson, nr) 629 if err != nil { 630 return err 631 } 632 defer b.bufSlabPredictor.put(jsonBuf) 633 if err := b.args[0].VecEvalJSON(b.ctx, input, jsonBuf); err != nil { 634 return err 635 } 636 pathBuf, err := b.bufSlabPredictor.get(types.ETString, nr) 637 if err != nil { 638 return err 639 } 640 defer b.bufSlabPredictor.put(pathBuf) 641 if err := b.args[1].VecEvalString(b.ctx, input, pathBuf); err != nil { 642 return err 643 } 644 645 result.ReserveJSON(nr) 646 for i := 0; i < nr; i++ { 647 if jsonBuf.IsNull(i) || pathBuf.IsNull(i) { 648 result.AppendNull() 649 continue 650 } 651 652 pathExpr, err := json.ParseJSONPathExpr(pathBuf.GetString(i)) 653 if err != nil { 654 return err 655 } 656 if pathExpr.ContainsAnyAsterisk() { 657 return json.ErrInvalidJSONPathWildcard 658 } 659 660 jsonItem := jsonBuf.GetJSON(i) 661 if jsonItem.TypeCode != json.TypeCodeObject { 662 result.AppendNull() 663 continue 664 } 665 666 res, exists := jsonItem.Extract([]json.PathExpression{pathExpr}) 667 if !exists || res.TypeCode != json.TypeCodeObject { 668 result.AppendNull() 669 continue 670 } 671 result.AppendJSON(res.GetKeys()) 672 } 673 return nil 674 } 675 676 func (b *builtinJSONLengthSig) vectorized() bool { 677 return true 678 } 679 680 func (b *builtinJSONLengthSig) vecEvalInt(input *chunk.Chunk, result *chunk.DeferredCauset) error { 681 nr := input.NumEvents() 682 683 jsonBuf, err := b.bufSlabPredictor.get(types.ETJson, nr) 684 if err != nil { 685 return err 686 } 687 defer b.bufSlabPredictor.put(jsonBuf) 688 if err := b.args[0].VecEvalJSON(b.ctx, input, jsonBuf); err != nil { 689 return err 690 } 691 result.ResizeInt64(nr, false) 692 resI64s := result.Int64s() 693 694 if len(b.args) == 2 { 695 pathBuf, err := b.bufSlabPredictor.get(types.ETString, nr) 696 if err != nil { 697 return err 698 } 699 defer b.bufSlabPredictor.put(pathBuf) 700 if err := b.args[1].VecEvalString(b.ctx, input, pathBuf); err != nil { 701 return err 702 } 703 704 result.MergeNulls(jsonBuf) 705 for i := 0; i < nr; i++ { 706 if result.IsNull(i) { 707 continue 708 } 709 jsonItem := jsonBuf.GetJSON(i) 710 711 if jsonItem.TypeCode != json.TypeCodeObject && jsonItem.TypeCode != json.TypeCodeArray { 712 resI64s[i] = 1 713 continue 714 } 715 716 if pathBuf.IsNull(i) { 717 result.SetNull(i, true) 718 continue 719 } 720 721 pathExpr, err := json.ParseJSONPathExpr(pathBuf.GetString(i)) 722 if err != nil { 723 return err 724 } 725 if pathExpr.ContainsAnyAsterisk() { 726 return json.ErrInvalidJSONPathWildcard 727 } 728 729 obj, exists := jsonItem.Extract([]json.PathExpression{pathExpr}) 730 if !exists { 731 result.SetNull(i, true) 732 continue 733 } 734 if obj.TypeCode != json.TypeCodeObject && obj.TypeCode != json.TypeCodeArray { 735 resI64s[i] = 1 736 continue 737 } 738 resI64s[i] = int64(obj.GetElemCount()) 739 } 740 } else { 741 result.MergeNulls(jsonBuf) 742 for i := 0; i < nr; i++ { 743 if result.IsNull(i) { 744 continue 745 } 746 747 jsonItem := jsonBuf.GetJSON(i) 748 if jsonItem.TypeCode != json.TypeCodeObject && jsonItem.TypeCode != json.TypeCodeArray { 749 resI64s[i] = 1 750 continue 751 } 752 resI64s[i] = int64(jsonItem.GetElemCount()) 753 } 754 } 755 756 return nil 757 } 758 759 func (b *builtinJSONTypeSig) vectorized() bool { 760 return true 761 } 762 763 func (b *builtinJSONTypeSig) vecEvalString(input *chunk.Chunk, result *chunk.DeferredCauset) error { 764 n := input.NumEvents() 765 buf, err := b.bufSlabPredictor.get(types.ETJson, n) 766 if err != nil { 767 return err 768 } 769 defer b.bufSlabPredictor.put(buf) 770 if err := b.args[0].VecEvalJSON(b.ctx, input, buf); err != nil { 771 return err 772 } 773 774 result.ReserveString(n) 775 for i := 0; i < n; i++ { 776 if buf.IsNull(i) { 777 result.AppendNull() 778 continue 779 } 780 result.AppendString(buf.GetJSON(i).Type()) 781 } 782 return nil 783 } 784 785 func (b *builtinJSONExtractSig) vectorized() bool { 786 return true 787 } 788 789 func (b *builtinJSONExtractSig) vecEvalJSON(input *chunk.Chunk, result *chunk.DeferredCauset) error { 790 var err error 791 792 nr := input.NumEvents() 793 jsonBuf, err := b.bufSlabPredictor.get(types.ETJson, nr) 794 if err != nil { 795 return err 796 } 797 defer b.bufSlabPredictor.put(jsonBuf) 798 if err = b.args[0].VecEvalJSON(b.ctx, input, jsonBuf); err != nil { 799 return err 800 } 801 802 pathArgs := b.args[1:] 803 pathBuffers := make([]*chunk.DeferredCauset, len(pathArgs)) 804 for k := 0; k < len(pathArgs); k++ { 805 if pathBuffers[k], err = b.bufSlabPredictor.get(types.ETString, nr); err != nil { 806 return err 807 } 808 defer func(buf *chunk.DeferredCauset) { 809 b.bufSlabPredictor.put(buf) 810 }(pathBuffers[k]) 811 812 if err := pathArgs[k].VecEvalString(b.ctx, input, pathBuffers[k]); err != nil { 813 return err 814 } 815 } 816 817 result.ReserveJSON(nr) 818 for i := 0; i < nr; i++ { 819 if jsonBuf.IsNull(i) { 820 result.AppendNull() 821 continue 822 } 823 jsonItem := jsonBuf.GetJSON(i) 824 825 pathExprs := make([]json.PathExpression, len(pathBuffers)) 826 hasNullPath := false 827 for k, pathBuf := range pathBuffers { 828 if pathBuf.IsNull(i) { 829 hasNullPath = true 830 break 831 } 832 if pathExprs[k], err = json.ParseJSONPathExpr(pathBuf.GetString(i)); err != nil { 833 return err 834 } 835 } 836 if hasNullPath { 837 result.AppendNull() 838 continue 839 } 840 841 var found bool 842 if jsonItem, found = jsonItem.Extract(pathExprs); !found { 843 result.AppendNull() 844 continue 845 } 846 result.AppendJSON(jsonItem) 847 } 848 849 return nil 850 } 851 852 func (b *builtinJSONRemoveSig) vectorized() bool { 853 return true 854 } 855 856 func (b *builtinJSONRemoveSig) vecEvalJSON(input *chunk.Chunk, result *chunk.DeferredCauset) error { 857 nr := input.NumEvents() 858 jsonBuf, err := b.bufSlabPredictor.get(types.ETJson, nr) 859 if err != nil { 860 return err 861 } 862 defer b.bufSlabPredictor.put(jsonBuf) 863 if err := b.args[0].VecEvalJSON(b.ctx, input, jsonBuf); err != nil { 864 return err 865 } 866 867 strBufs := make([]*chunk.DeferredCauset, len(b.args)-1) 868 for i := 1; i < len(b.args); i++ { 869 strBufs[i-1], err = b.bufSlabPredictor.get(types.ETString, nr) 870 if err != nil { 871 return err 872 } 873 defer b.bufSlabPredictor.put(strBufs[i-1]) 874 if err := b.args[i].VecEvalString(b.ctx, input, strBufs[i-1]); err != nil { 875 return err 876 } 877 } 878 879 result.ReserveJSON(nr) 880 for i := 0; i < nr; i++ { 881 if jsonBuf.IsNull(i) { 882 result.AppendNull() 883 continue 884 } 885 886 pathExprs := make([]json.PathExpression, 0, len(b.args)-1) 887 var pathExpr json.PathExpression 888 isNull := false 889 890 for j := 1; j < len(b.args); j++ { 891 if strBufs[j-1].IsNull(i) { 892 isNull = true 893 break 894 } 895 pathExpr, err = json.ParseJSONPathExpr(strBufs[j-1].GetString(i)) 896 if err != nil { 897 return err 898 } 899 pathExprs = append(pathExprs, pathExpr) 900 } 901 902 if isNull { 903 result.AppendNull() 904 } else { 905 res, err := jsonBuf.GetJSON(i).Remove(pathExprs) 906 if err != nil { 907 return err 908 } 909 result.AppendJSON(res) 910 } 911 } 912 return nil 913 } 914 915 func (b *builtinJSONMergeSig) vectorized() bool { 916 return true 917 } 918 919 func (b *builtinJSONMergeSig) vecEvalJSON(input *chunk.Chunk, result *chunk.DeferredCauset) error { 920 nr := input.NumEvents() 921 argBuffers := make([]*chunk.DeferredCauset, len(b.args)) 922 var err error 923 for i, arg := range b.args { 924 if argBuffers[i], err = b.bufSlabPredictor.get(types.ETJson, nr); err != nil { 925 return err 926 } 927 defer func(buf *chunk.DeferredCauset) { 928 b.bufSlabPredictor.put(buf) 929 }(argBuffers[i]) 930 931 if err := arg.VecEvalJSON(b.ctx, input, argBuffers[i]); err != nil { 932 return err 933 } 934 } 935 936 jsonValues := make([][]json.BinaryJSON, nr) 937 938 for i := 0; i < nr; i++ { 939 isNullFlag := false 940 for j := 0; j < len(b.args); j++ { 941 isNullFlag = isNullFlag || argBuffers[j].IsNull(i) 942 } 943 if isNullFlag { 944 jsonValues[i] = nil 945 } else { 946 jsonValues[i] = make([]json.BinaryJSON, 0, len(b.args)) 947 } 948 } 949 for i := 0; i < len(b.args); i++ { 950 for j := 0; j < nr; j++ { 951 if jsonValues[j] == nil { 952 continue 953 } 954 jsonValues[j] = append(jsonValues[j], argBuffers[i].GetJSON(j)) 955 } 956 } 957 958 result.ReserveJSON(nr) 959 for i := 0; i < nr; i++ { 960 if jsonValues[i] == nil { 961 result.AppendNull() 962 continue 963 } 964 result.AppendJSON(json.MergeBinary(jsonValues[i])) 965 } 966 967 if b.pbCode == fidelpb.ScalarFuncSig_JsonMergeSig { 968 for i := 0; i < nr; i++ { 969 if result.IsNull(i) { 970 continue 971 } 972 b.ctx.GetStochastikVars().StmtCtx.AppendWarning(errDeprecatedSyntaxNoRememristed.GenWithStackByArgs("JSON_MERGE")) 973 } 974 } 975 976 return nil 977 } 978 979 func (b *builtinJSONContainsPathSig) vectorized() bool { 980 return true 981 } 982 983 func (b *builtinJSONContainsPathSig) vecEvalInt(input *chunk.Chunk, result *chunk.DeferredCauset) error { 984 n := input.NumEvents() 985 jsonBuf, err := b.bufSlabPredictor.get(types.ETJson, n) 986 if err != nil { 987 return err 988 } 989 defer b.bufSlabPredictor.put(jsonBuf) 990 if err := b.args[0].VecEvalJSON(b.ctx, input, jsonBuf); err != nil { 991 return err 992 } 993 typeBuf, err := b.bufSlabPredictor.get(types.ETString, n) 994 if err != nil { 995 return err 996 } 997 defer b.bufSlabPredictor.put(typeBuf) 998 if err := b.args[1].VecEvalString(b.ctx, input, typeBuf); err != nil { 999 return err 1000 } 1001 pathBufs := make([]*chunk.DeferredCauset, len(b.args)-2) 1002 defer func() { 1003 for i := 0; i < len(pathBufs); i++ { 1004 if pathBufs[i] != nil { 1005 b.bufSlabPredictor.put(pathBufs[i]) 1006 } 1007 } 1008 }() 1009 for i := 0; i < len(pathBufs); i++ { 1010 pathBuf, err := b.bufSlabPredictor.get(types.ETString, n) 1011 if err != nil { 1012 return err 1013 } 1014 pathBufs[i] = pathBuf 1015 if err := b.args[2+i].VecEvalString(b.ctx, input, pathBuf); err != nil { 1016 return err 1017 } 1018 } 1019 1020 result.ResizeInt64(n, false) 1021 result.MergeNulls(jsonBuf, typeBuf) 1022 i64s := result.Int64s() 1023 for i := 0; i < n; i++ { 1024 if result.IsNull(i) { 1025 continue 1026 } 1027 containType := strings.ToLower(typeBuf.GetString(i)) 1028 if containType != json.ContainsPathAll && containType != json.ContainsPathOne { 1029 return json.ErrInvalidJSONContainsPathType 1030 } 1031 obj := jsonBuf.GetJSON(i) 1032 contains := int64(1) 1033 var pathExpr json.PathExpression 1034 for j := 0; j < len(pathBufs); j++ { 1035 if pathBufs[j].IsNull(i) { 1036 result.SetNull(i, true) 1037 contains = -1 1038 break 1039 } 1040 path := pathBufs[j].GetString(i) 1041 if pathExpr, err = json.ParseJSONPathExpr(path); err != nil { 1042 return err 1043 } 1044 _, exists := obj.Extract([]json.PathExpression{pathExpr}) 1045 if exists && containType == json.ContainsPathOne { 1046 contains = 1 1047 break 1048 } else if !exists && containType == json.ContainsPathOne { 1049 contains = 0 1050 } else if !exists && containType == json.ContainsPathAll { 1051 contains = 0 1052 break 1053 } 1054 } 1055 if contains >= 0 { 1056 i64s[i] = contains 1057 } 1058 } 1059 return nil 1060 } 1061 1062 func (b *builtinJSONArrayAppendSig) vectorized() bool { 1063 return true 1064 } 1065 1066 func (b *builtinJSONArrayAppendSig) vecEvalJSON(input *chunk.Chunk, result *chunk.DeferredCauset) error { 1067 n := input.NumEvents() 1068 m := (len(b.args) - 1) / 2 1069 1070 jsonBufs, err := b.bufSlabPredictor.get(types.ETJson, n) 1071 if err != nil { 1072 return err 1073 } 1074 defer b.bufSlabPredictor.put(jsonBufs) 1075 if err := b.args[0].VecEvalJSON(b.ctx, input, jsonBufs); err != nil { 1076 return err 1077 } 1078 1079 pathBufs := make([]*chunk.DeferredCauset, 0, m) 1080 valBufs := make([]*chunk.DeferredCauset, 0, m) 1081 defer func() { 1082 for _, buf := range pathBufs { 1083 b.bufSlabPredictor.put(buf) 1084 } 1085 for _, buf := range valBufs { 1086 b.bufSlabPredictor.put(buf) 1087 } 1088 }() 1089 for i := 1; i < len(b.args)-1; i += 2 { 1090 pathBuf, err := b.bufSlabPredictor.get(types.ETString, n) 1091 if err != nil { 1092 return err 1093 } 1094 pathBufs = append(pathBufs, pathBuf) 1095 if err := b.args[i].VecEvalString(b.ctx, input, pathBuf); err != nil { 1096 return err 1097 } 1098 valBuf, err := b.bufSlabPredictor.get(types.ETJson, n) 1099 if err != nil { 1100 return err 1101 } 1102 if err := b.args[i+1].VecEvalJSON(b.ctx, input, valBuf); err != nil { 1103 return err 1104 } 1105 valBufs = append(valBufs, valBuf) 1106 } 1107 1108 result.ReserveJSON(n) 1109 for i := 0; i < n; i++ { 1110 if jsonBufs.IsNull(i) { 1111 result.AppendNull() 1112 continue 1113 } 1114 res := jsonBufs.GetJSON(i) 1115 isNull := false 1116 for j := 0; j < m; j++ { 1117 if pathBufs[j].IsNull(i) { 1118 isNull = true 1119 break 1120 } 1121 s := pathBufs[j].GetString(i) 1122 v, vNull := json.BinaryJSON{}, valBufs[j].IsNull(i) 1123 if !vNull { 1124 v = valBufs[j].GetJSON(i) 1125 } 1126 if vNull { 1127 v = json.CreateBinary(nil) 1128 } 1129 res, isNull, err = b.appendJSONArray(res, s, v) 1130 if err != nil { 1131 return err 1132 } 1133 if isNull { 1134 break 1135 } 1136 } 1137 if isNull { 1138 result.AppendNull() 1139 } else { 1140 result.AppendJSON(res) 1141 } 1142 } 1143 return nil 1144 } 1145 1146 func (b *builtinJSONUnquoteSig) vectorized() bool { 1147 return true 1148 } 1149 1150 func (b *builtinJSONUnquoteSig) vecEvalString(input *chunk.Chunk, result *chunk.DeferredCauset) error { 1151 n := input.NumEvents() 1152 buf, err := b.bufSlabPredictor.get(types.ETString, n) 1153 if err != nil { 1154 return err 1155 } 1156 defer b.bufSlabPredictor.put(buf) 1157 if err := b.args[0].VecEvalString(b.ctx, input, buf); err != nil { 1158 return err 1159 } 1160 1161 result.ReserveString(n) 1162 for i := 0; i < n; i++ { 1163 if buf.IsNull(i) { 1164 result.AppendNull() 1165 continue 1166 } 1167 str, err := json.UnquoteString(buf.GetString(i)) 1168 if err != nil { 1169 return err 1170 } 1171 result.AppendString(str) 1172 } 1173 return nil 1174 }