github.com/apache/arrow/go/v14@v14.0.1/arrow/compute/scalar_compare_test.go (about) 1 // Licensed to the Apache Software Foundation (ASF) under one 2 // or more contributor license agreements. See the NOTICE file 3 // distributed with this work for additional information 4 // regarding copyright ownership. The ASF licenses this file 5 // to you under the Apache License, Version 2.0 (the 6 // "License"); you may not use this file except in compliance 7 // with the License. You may obtain a copy of the License at 8 // 9 // http://www.apache.org/licenses/LICENSE-2.0 10 // 11 // Unless required by applicable law or agreed to in writing, software 12 // distributed under the License is distributed on an "AS IS" BASIS, 13 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 // See the License for the specific language governing permissions and 15 // limitations under the License. 16 17 //go:build go1.18 18 19 package compute_test 20 21 import ( 22 "context" 23 "fmt" 24 "strings" 25 "testing" 26 27 "github.com/apache/arrow/go/v14/arrow" 28 "github.com/apache/arrow/go/v14/arrow/array" 29 "github.com/apache/arrow/go/v14/arrow/bitutil" 30 "github.com/apache/arrow/go/v14/arrow/compute" 31 "github.com/apache/arrow/go/v14/arrow/compute/exec" 32 "github.com/apache/arrow/go/v14/arrow/compute/internal/kernels" 33 "github.com/apache/arrow/go/v14/arrow/internal/testing/gen" 34 "github.com/apache/arrow/go/v14/arrow/memory" 35 "github.com/apache/arrow/go/v14/arrow/scalar" 36 "github.com/stretchr/testify/assert" 37 "github.com/stretchr/testify/require" 38 "github.com/stretchr/testify/suite" 39 ) 40 41 type CompareSuite struct { 42 BinaryFuncTestSuite 43 } 44 45 func (c *CompareSuite) validateCompareDatum(op kernels.CompareOperator, lhs, rhs, expected compute.Datum) { 46 result, err := compute.CallFunction(c.ctx, op.String(), nil, lhs, rhs) 47 c.Require().NoError(err) 48 defer result.Release() 49 50 assertDatumsEqual(c.T(), expected, result, nil, nil) 51 } 52 53 func (c *CompareSuite) validateCompare(op kernels.CompareOperator, dt arrow.DataType, lhsStr, rhsStr, expStr string) { 54 lhs, _, err := array.FromJSON(c.mem, dt, strings.NewReader(lhsStr), array.WithUseNumber()) 55 c.Require().NoError(err) 56 rhs, _, err := array.FromJSON(c.mem, dt, strings.NewReader(rhsStr), array.WithUseNumber()) 57 c.Require().NoError(err) 58 exp, _, err := array.FromJSON(c.mem, arrow.FixedWidthTypes.Boolean, strings.NewReader(expStr), array.WithUseNumber()) 59 c.Require().NoError(err) 60 defer func() { 61 lhs.Release() 62 rhs.Release() 63 exp.Release() 64 }() 65 c.validateCompareDatum(op, &compute.ArrayDatum{lhs.Data()}, &compute.ArrayDatum{rhs.Data()}, &compute.ArrayDatum{exp.Data()}) 66 } 67 68 func (c *CompareSuite) validateCompareArrScalar(op kernels.CompareOperator, dt arrow.DataType, lhsStr string, rhs compute.Datum, expStr string) { 69 lhs, _, err := array.FromJSON(c.mem, dt, strings.NewReader(lhsStr), array.WithUseNumber()) 70 c.Require().NoError(err) 71 exp, _, err := array.FromJSON(c.mem, arrow.FixedWidthTypes.Boolean, strings.NewReader(expStr), array.WithUseNumber()) 72 c.Require().NoError(err) 73 defer func() { 74 lhs.Release() 75 exp.Release() 76 }() 77 c.validateCompareDatum(op, &compute.ArrayDatum{lhs.Data()}, rhs, &compute.ArrayDatum{exp.Data()}) 78 } 79 80 func (c *CompareSuite) validateCompareScalarArr(op kernels.CompareOperator, dt arrow.DataType, lhs compute.Datum, rhsStr string, expStr string) { 81 rhs, _, err := array.FromJSON(c.mem, dt, strings.NewReader(rhsStr), array.WithUseNumber()) 82 c.Require().NoError(err) 83 exp, _, err := array.FromJSON(c.mem, arrow.FixedWidthTypes.Boolean, strings.NewReader(expStr), array.WithUseNumber()) 84 c.Require().NoError(err) 85 defer func() { 86 rhs.Release() 87 exp.Release() 88 }() 89 c.validateCompareDatum(op, lhs, &compute.ArrayDatum{rhs.Data()}, &compute.ArrayDatum{exp.Data()}) 90 } 91 92 func slowCompare[T exec.NumericTypes | string](op kernels.CompareOperator, lhs, rhs T) bool { 93 switch op { 94 case kernels.CmpEQ: 95 return lhs == rhs 96 case kernels.CmpNE: 97 return lhs != rhs 98 case kernels.CmpLT: 99 return lhs < rhs 100 case kernels.CmpLE: 101 return lhs <= rhs 102 case kernels.CmpGT: 103 return lhs > rhs 104 case kernels.CmpGE: 105 return lhs >= rhs 106 default: 107 return false 108 } 109 } 110 111 // func simpleScalarArrayCompare[T exec.NumericTypes](mem memory.Allocator, op kernels.CompareOperator, lhs, rhs compute.Datum) compute.Datum { 112 // var ( 113 // swap = lhs.Kind() == compute.KindArray 114 // span exec.ArraySpan 115 // itr exec.ArrayIter[T] 116 // value T 117 // ) 118 119 // if swap { 120 // span.SetMembers(lhs.(*compute.ArrayDatum).Value) 121 // itr = exec.NewPrimitiveIter[T](&span) 122 // value = kernels.UnboxScalar[T](rhs.(*compute.ScalarDatum).Value.(scalar.PrimitiveScalar)) 123 // } else { 124 // span.SetMembers(rhs.(*compute.ArrayDatum).Value) 125 // itr = exec.NewPrimitiveIter[T](&span) 126 // value = kernels.UnboxScalar[T](lhs.(*compute.ScalarDatum).Value.(scalar.PrimitiveScalar)) 127 // } 128 129 // bitmap := make([]bool, span.Len) 130 // for i := 0; i < int(span.Len); i++ { 131 // if swap { 132 // bitmap[i] = slowCompare(op, itr.Next(), value) 133 // } else { 134 // bitmap[i] = slowCompare(op, value, itr.Next()) 135 // } 136 // } 137 138 // var result arrow.Array 139 // if span.Nulls == 0 { 140 // result = exec.ArrayFromSlice(mem, bitmap) 141 // } else { 142 // nullBitmap := make([]bool, span.Len) 143 // rdr := bitutil.NewBitmapReader(span.Buffers[0].Buf, int(span.Offset), int(span.Len)) 144 // for i := 0; i < int(span.Len); i++ { 145 // nullBitmap[i] = rdr.Set() 146 // rdr.Next() 147 // } 148 // bldr := array.NewBooleanBuilder(mem) 149 // defer bldr.Release() 150 151 // bldr.AppendValues(bitmap, nullBitmap) 152 // result = bldr.NewArray() 153 // } 154 155 // defer result.Release() 156 // return compute.NewDatum(result) 157 // } 158 159 func simpleScalarArrayCompareString(mem memory.Allocator, op kernels.CompareOperator, lhs, rhs compute.Datum) compute.Datum { 160 var ( 161 swap = lhs.Kind() == compute.KindArray 162 value string 163 arr *array.String 164 ) 165 166 if swap { 167 arr = lhs.(*compute.ArrayDatum).MakeArray().(*array.String) 168 defer arr.Release() 169 value = string(rhs.(*compute.ScalarDatum).Value.(*scalar.String).Data()) 170 } else { 171 arr = rhs.(*compute.ArrayDatum).MakeArray().(*array.String) 172 defer arr.Release() 173 value = string(lhs.(*compute.ScalarDatum).Value.(*scalar.String).Data()) 174 } 175 176 bitmap := make([]bool, arr.Len()) 177 for i := 0; i < arr.Len(); i++ { 178 if swap { 179 bitmap[i] = slowCompare(op, arr.Value(i), value) 180 } else { 181 bitmap[i] = slowCompare(op, value, arr.Value(i)) 182 } 183 } 184 185 var result arrow.Array 186 if arr.NullN() == 0 { 187 result = exec.ArrayFromSlice(mem, bitmap) 188 } else { 189 nullBitmap := make([]bool, arr.Len()) 190 rdr := bitutil.NewBitmapReader(arr.NullBitmapBytes(), arr.Offset(), arr.Len()) 191 for i := 0; i < arr.Len(); i++ { 192 nullBitmap[i] = rdr.Set() 193 rdr.Next() 194 } 195 bldr := array.NewBooleanBuilder(mem) 196 defer bldr.Release() 197 198 bldr.AppendValues(bitmap, nullBitmap) 199 result = bldr.NewArray() 200 } 201 202 defer result.Release() 203 return compute.NewDatum(result) 204 } 205 206 func nullBitmapFromArrays(lhs, rhs arrow.Array) []bool { 207 nullBitmap := make([]bool, lhs.Len()) 208 209 left := func(i int) bool { 210 if lhs.NullN() == 0 { 211 return true 212 } 213 return lhs.IsValid(i) 214 } 215 216 right := func(i int) bool { 217 if rhs.NullN() == 0 { 218 return true 219 } 220 return rhs.IsValid(i) 221 } 222 223 for i := 0; i < lhs.Len(); i++ { 224 nullBitmap[i] = left(i) && right(i) 225 } 226 return nullBitmap 227 } 228 229 type valuer[T any] interface { 230 Value(int) T 231 } 232 233 func simpleArrArrCompare[T exec.NumericTypes | string](mem memory.Allocator, op kernels.CompareOperator, lhs, rhs compute.Datum) compute.Datum { 234 var ( 235 lArr = lhs.(*compute.ArrayDatum).MakeArray() 236 rArr = rhs.(*compute.ArrayDatum).MakeArray() 237 length = lArr.Len() 238 bitmap = make([]bool, length) 239 240 lvals = lArr.(valuer[T]) 241 rvals = rArr.(valuer[T]) 242 ) 243 defer lArr.Release() 244 defer rArr.Release() 245 246 for i := 0; i < length; i++ { 247 bitmap[i] = slowCompare(op, lvals.Value(i), rvals.Value(i)) 248 } 249 250 var result arrow.Array 251 if lArr.NullN() == 0 && rArr.NullN() == 0 { 252 result = exec.ArrayFromSlice(mem, bitmap) 253 } else { 254 nullBitmap := nullBitmapFromArrays(lArr, rArr) 255 bldr := array.NewBooleanBuilder(mem) 256 defer bldr.Release() 257 258 bldr.AppendValues(bitmap, nullBitmap) 259 result = bldr.NewArray() 260 } 261 262 defer result.Release() 263 return compute.NewDatum(result) 264 } 265 266 type NumericCompareSuite[T exec.NumericTypes] struct { 267 CompareSuite 268 } 269 270 // func (n *NumericCompareSuite[T]) validateCompareComputed(op kernels.CompareOperator, lhs, rhs compute.Datum) { 271 // var expected compute.Datum 272 273 // hasScalar := lhs.Kind() == compute.KindScalar || rhs.Kind() == compute.KindScalar 274 // if hasScalar { 275 // expected = simpleScalarArrayCompare[T](n.mem, op, lhs, rhs) 276 // } else { 277 // expected = simpleArrArrCompare[T](n.mem, op, lhs, rhs) 278 // } 279 280 // defer expected.Release() 281 // n.CompareSuite.validateCompareDatum(op, lhs, rhs, expected) 282 // } 283 284 func (n *NumericCompareSuite[T]) TestSimpleCompareArrayScalar() { 285 dt := exec.GetDataType[T]() 286 one := compute.NewDatum(scalar.MakeScalar(T(1))) 287 288 n.Run(dt.String(), func() { 289 op := kernels.CmpEQ 290 n.validateCompareArrScalar(op, dt, `[]`, one, `[]`) 291 n.validateCompareArrScalar(op, dt, `[null]`, one, `[null]`) 292 n.validateCompareArrScalar(op, dt, `[0, 0, 1, 1, 2, 2]`, one, 293 `[false, false, true, true, false, false]`) 294 n.validateCompareArrScalar(op, dt, `[0, 1, 2, 3, 4, 5]`, one, 295 `[false, true, false, false, false, false]`) 296 n.validateCompareArrScalar(op, dt, `[5, 4, 3, 2, 1, 0]`, one, 297 `[false, false, false, false, true, false]`) 298 n.validateCompareArrScalar(op, dt, `[null, 0, 1, 1]`, one, 299 `[null, false, true, true]`) 300 301 op = kernels.CmpNE 302 n.validateCompareArrScalar(op, dt, `[]`, one, `[]`) 303 n.validateCompareArrScalar(op, dt, `[null]`, one, `[null]`) 304 n.validateCompareArrScalar(op, dt, `[0, 0, 1, 1, 2, 2]`, one, 305 `[true, true, false, false, true, true]`) 306 n.validateCompareArrScalar(op, dt, `[0, 1, 2, 3, 4, 5]`, one, 307 `[true, false, true, true, true, true]`) 308 n.validateCompareArrScalar(op, dt, `[5, 4, 3, 2, 1, 0]`, one, 309 `[true, true, true, true, false, true]`) 310 n.validateCompareArrScalar(op, dt, `[null, 0, 1, 1]`, one, 311 `[null, true, false, false]`) 312 313 op = kernels.CmpGT 314 n.validateCompareArrScalar(op, dt, `[]`, one, `[]`) 315 n.validateCompareArrScalar(op, dt, `[null]`, one, `[null]`) 316 n.validateCompareArrScalar(op, dt, `[0, 0, 1, 1, 2, 2]`, one, 317 `[false, false, false, false, true, true]`) 318 n.validateCompareArrScalar(op, dt, `[0, 1, 2, 3, 4, 5]`, one, 319 `[false, false, true, true, true, true]`) 320 n.validateCompareArrScalar(op, dt, `[4, 5, 6, 7, 8, 9]`, one, 321 `[true, true, true, true, true, true]`) 322 n.validateCompareArrScalar(op, dt, `[null, 0, 1, 1]`, one, 323 `[null, false, false, false]`) 324 325 op = kernels.CmpGE 326 n.validateCompareArrScalar(op, dt, `[]`, one, `[]`) 327 n.validateCompareArrScalar(op, dt, `[null]`, one, `[null]`) 328 n.validateCompareArrScalar(op, dt, `[0, 0, 1, 1, 2, 2]`, one, 329 `[false, false, true, true, true, true]`) 330 n.validateCompareArrScalar(op, dt, `[0, 1, 2, 3, 4, 5]`, one, 331 `[false, true, true, true, true, true]`) 332 n.validateCompareArrScalar(op, dt, `[4, 5, 6, 7, 8, 9]`, one, 333 `[true, true, true, true, true, true]`) 334 n.validateCompareArrScalar(op, dt, `[null, 0, 1, 1]`, one, 335 `[null, false, true, true]`) 336 337 op = kernels.CmpLT 338 n.validateCompareArrScalar(op, dt, `[]`, one, `[]`) 339 n.validateCompareArrScalar(op, dt, `[null]`, one, `[null]`) 340 n.validateCompareArrScalar(op, dt, `[0, 0, 1, 1, 2, 2]`, one, 341 `[true, true, false, false, false, false]`) 342 n.validateCompareArrScalar(op, dt, `[0, 1, 2, 3, 4, 5]`, one, 343 `[true, false, false, false, false, false]`) 344 n.validateCompareArrScalar(op, dt, `[4, 5, 6, 7, 8, 9]`, one, 345 `[false, false, false, false, false, false]`) 346 n.validateCompareArrScalar(op, dt, `[null, 0, 1, 1]`, one, 347 `[null, true, false, false]`) 348 349 op = kernels.CmpLE 350 n.validateCompareArrScalar(op, dt, `[]`, one, `[]`) 351 n.validateCompareArrScalar(op, dt, `[null]`, one, `[null]`) 352 n.validateCompareArrScalar(op, dt, `[0, 0, 1, 1, 2, 2]`, one, 353 `[true, true, true, true, false, false]`) 354 n.validateCompareArrScalar(op, dt, `[0, 1, 2, 3, 4, 5]`, one, 355 `[true, true, false, false, false, false]`) 356 n.validateCompareArrScalar(op, dt, `[4, 5, 6, 7, 8, 9]`, one, 357 `[false, false, false, false, false, false]`) 358 n.validateCompareArrScalar(op, dt, `[null, 0, 1, 1]`, one, 359 `[null, true, true, true]`) 360 }) 361 } 362 363 func (n *NumericCompareSuite[T]) TestSimpleCompareScalarArray() { 364 dt := exec.GetDataType[T]() 365 one := compute.NewDatum(scalar.MakeScalar(T(1))) 366 367 n.Run(dt.String(), func() { 368 op := kernels.CmpEQ 369 n.validateCompareScalarArr(op, dt, one, `[]`, `[]`) 370 n.validateCompareScalarArr(op, dt, one, `[null]`, `[null]`) 371 n.validateCompareScalarArr(op, dt, one, `[0, 0, 1, 1, 2, 2]`, 372 `[false, false, true, true, false, false]`) 373 n.validateCompareScalarArr(op, dt, one, `[0, 1, 2, 3, 4, 5]`, 374 `[false, true, false, false, false, false]`) 375 n.validateCompareScalarArr(op, dt, one, `[5, 4, 3, 2, 1, 0]`, 376 `[false, false, false, false, true, false]`) 377 n.validateCompareScalarArr(op, dt, one, `[null, 0, 1, 1]`, 378 `[null, false, true, true]`) 379 380 op = kernels.CmpNE 381 n.validateCompareScalarArr(op, dt, one, `[]`, `[]`) 382 n.validateCompareScalarArr(op, dt, one, `[null]`, `[null]`) 383 n.validateCompareScalarArr(op, dt, one, `[0, 0, 1, 1, 2, 2]`, 384 `[true, true, false, false, true, true]`) 385 n.validateCompareScalarArr(op, dt, one, `[0, 1, 2, 3, 4, 5]`, 386 `[true, false, true, true, true, true]`) 387 n.validateCompareScalarArr(op, dt, one, `[5, 4, 3, 2, 1, 0]`, 388 `[true, true, true, true, false, true]`) 389 n.validateCompareScalarArr(op, dt, one, `[null, 0, 1, 1]`, 390 `[null, true, false, false]`) 391 392 op = kernels.CmpGT 393 n.validateCompareScalarArr(op, dt, one, `[]`, `[]`) 394 n.validateCompareScalarArr(op, dt, one, `[null]`, `[null]`) 395 n.validateCompareScalarArr(op, dt, one, `[0, 0, 1, 1, 2, 2]`, 396 `[true, true, false, false, false, false]`) 397 n.validateCompareScalarArr(op, dt, one, `[0, 1, 2, 3, 4, 5]`, 398 `[true, false, false, false, false, false]`) 399 n.validateCompareScalarArr(op, dt, one, `[4, 5, 6, 7, 8, 9]`, 400 `[false, false, false, false, false, false]`) 401 n.validateCompareScalarArr(op, dt, one, `[null, 0, 1, 1]`, 402 `[null, true, false, false]`) 403 404 op = kernels.CmpGE 405 n.validateCompareScalarArr(op, dt, one, `[]`, `[]`) 406 n.validateCompareScalarArr(op, dt, one, `[null]`, `[null]`) 407 n.validateCompareScalarArr(op, dt, one, `[0, 0, 1, 1, 2, 2]`, 408 `[true, true, true, true, false, false]`) 409 n.validateCompareScalarArr(op, dt, one, `[0, 1, 2, 3, 4, 5]`, 410 `[true, true, false, false, false, false]`) 411 n.validateCompareScalarArr(op, dt, one, `[4, 5, 6, 7, 8, 9]`, 412 `[false, false, false, false, false, false]`) 413 n.validateCompareScalarArr(op, dt, one, `[null, 0, 1, 1]`, 414 `[null, true, true, true]`) 415 416 op = kernels.CmpLT 417 n.validateCompareScalarArr(op, dt, one, `[]`, `[]`) 418 n.validateCompareScalarArr(op, dt, one, `[null]`, `[null]`) 419 n.validateCompareScalarArr(op, dt, one, `[0, 0, 1, 1, 2, 2]`, 420 `[false, false, false, false, true, true]`) 421 n.validateCompareScalarArr(op, dt, one, `[0, 1, 2, 3, 4, 5]`, 422 `[false, false, true, true, true, true]`) 423 n.validateCompareScalarArr(op, dt, one, `[4, 5, 6, 7, 8, 9]`, 424 `[true, true, true, true, true, true]`) 425 n.validateCompareScalarArr(op, dt, one, `[null, 0, 1, 1]`, 426 `[null, false, false, false]`) 427 428 op = kernels.CmpLE 429 n.validateCompareScalarArr(op, dt, one, `[]`, `[]`) 430 n.validateCompareScalarArr(op, dt, one, `[null]`, `[null]`) 431 n.validateCompareScalarArr(op, dt, one, `[0, 0, 1, 1, 2, 2]`, 432 `[false, false, true, true, true, true]`) 433 n.validateCompareScalarArr(op, dt, one, `[0, 1, 2, 3, 4, 5]`, 434 `[false, true, true, true, true, true]`) 435 n.validateCompareScalarArr(op, dt, one, `[4, 5, 6, 7, 8, 9]`, 436 `[true, true, true, true, true, true]`) 437 n.validateCompareScalarArr(op, dt, one, `[null, 0, 1, 1]`, 438 `[null, false, true, true]`) 439 }) 440 } 441 442 func (n *NumericCompareSuite[T]) TestNullScalar() { 443 dt := exec.GetDataType[T]() 444 null := compute.NewDatum(scalar.MakeNullScalar(dt)) 445 446 n.Run(dt.String(), func() { 447 n.validateCompareArrScalar(kernels.CmpEQ, dt, `[]`, null, `[]`) 448 n.validateCompareScalarArr(kernels.CmpEQ, dt, null, `[]`, `[]`) 449 n.validateCompareArrScalar(kernels.CmpEQ, dt, `[null]`, null, `[null]`) 450 n.validateCompareScalarArr(kernels.CmpEQ, dt, null, `[null]`, `[null]`) 451 n.validateCompareScalarArr(kernels.CmpEQ, dt, null, `[1, 2, 3]`, `[null, null, null]`) 452 }) 453 } 454 455 func (n *NumericCompareSuite[T]) TestSimpleCompareArrArr() { 456 dt := exec.GetDataType[T]() 457 458 n.Run(dt.String(), func() { 459 n.validateCompare(kernels.CmpEQ, dt, `[]`, `[]`, `[]`) 460 n.validateCompare(kernels.CmpEQ, dt, `[null]`, `[null]`, `[null]`) 461 n.validateCompare(kernels.CmpEQ, dt, `[1]`, `[1]`, `[true]`) 462 n.validateCompare(kernels.CmpEQ, dt, `[1]`, `[2]`, `[false]`) 463 n.validateCompare(kernels.CmpEQ, dt, `[null]`, `[1]`, `[null]`) 464 n.validateCompare(kernels.CmpEQ, dt, `[1]`, `[null]`, `[null]`) 465 466 n.validateCompare(kernels.CmpLE, dt, `[1, 2, 3, 4, 5]`, `[2, 3, 4, 5, 6]`, `[true, true, true, true, true]`) 467 }) 468 } 469 470 type CompareTimestampSuite struct { 471 CompareSuite 472 } 473 474 func (c *CompareTimestampSuite) TestBasics() { 475 var ( 476 example1JSON = `["1970-01-01", "2000-02-29", "1900-02-28"]` 477 example2JSON = `["1970-01-02", "2000-02-01", "1900-02-28"]` 478 ) 479 480 checkCase := func(dt arrow.DataType, op kernels.CompareOperator, expected string) { 481 c.validateCompare(op, dt, example1JSON, example2JSON, expected) 482 } 483 484 seconds := arrow.FixedWidthTypes.Timestamp_s 485 millis := arrow.FixedWidthTypes.Timestamp_ms 486 micro := arrow.FixedWidthTypes.Timestamp_us 487 nano := arrow.FixedWidthTypes.Timestamp_ns 488 489 checkCase(seconds, kernels.CmpEQ, `[false, false, true]`) 490 checkCase(millis, kernels.CmpEQ, `[false, false, true]`) 491 checkCase(micro, kernels.CmpEQ, `[false, false, true]`) 492 checkCase(nano, kernels.CmpEQ, `[false, false, true]`) 493 494 checkCase(seconds, kernels.CmpNE, `[true, true, false]`) 495 checkCase(millis, kernels.CmpNE, `[true, true, false]`) 496 checkCase(micro, kernels.CmpNE, `[true, true, false]`) 497 checkCase(nano, kernels.CmpNE, `[true, true, false]`) 498 499 checkCase(seconds, kernels.CmpLT, `[true, false, false]`) 500 checkCase(seconds, kernels.CmpLE, `[true, false, true]`) 501 checkCase(seconds, kernels.CmpGT, `[false, true, false]`) 502 checkCase(seconds, kernels.CmpGE, `[false, true, true]`) 503 504 secondsUTC := &arrow.TimestampType{Unit: arrow.Second, TimeZone: "utc"} 505 checkCase(secondsUTC, kernels.CmpEQ, `[false, false, true]`) 506 } 507 508 func (c *CompareTimestampSuite) TestDiffParams() { 509 cases := []struct { 510 fn string 511 exp string 512 }{ 513 {"equal", `[false, false, true]`}, 514 {"not_equal", `[true, true, false]`}, 515 {"less", `[true, false, false]`}, 516 {"less_equal", `[true, false, true]`}, 517 {"greater", `[false, true, false]`}, 518 {"greater_equal", `[false, true, true]`}, 519 } 520 521 const lhsJSON = `["1970-01-01", "2000-02-29", "1900-02-28"]` 522 const rhsJSON = `["1970-01-02", "2000-02-01", "1900-02-28"]` 523 524 for _, op := range cases { 525 c.Run(op.fn, func() { 526 exp := c.getArr(arrow.FixedWidthTypes.Boolean, op.exp) 527 defer exp.Release() 528 529 expected := &compute.ArrayDatum{exp.Data()} 530 c.Run("diff units", func() { 531 lhs := c.getArr(&arrow.TimestampType{Unit: arrow.Second}, lhsJSON) 532 defer lhs.Release() 533 rhs := c.getArr(&arrow.TimestampType{Unit: arrow.Millisecond}, rhsJSON) 534 defer rhs.Release() 535 536 checkScalarBinary(c.T(), op.fn, &compute.ArrayDatum{lhs.Data()}, &compute.ArrayDatum{rhs.Data()}, expected, nil) 537 }) 538 c.Run("diff time zones", func() { 539 lhs := c.getArr(&arrow.TimestampType{Unit: arrow.Second, TimeZone: "America/New_York"}, lhsJSON) 540 defer lhs.Release() 541 rhs := c.getArr(&arrow.TimestampType{Unit: arrow.Second, TimeZone: "America/Phoenix"}, rhsJSON) 542 defer rhs.Release() 543 544 checkScalarBinary(c.T(), op.fn, &compute.ArrayDatum{lhs.Data()}, &compute.ArrayDatum{rhs.Data()}, expected, nil) 545 }) 546 c.Run("native to zoned", func() { 547 lhs := c.getArr(&arrow.TimestampType{Unit: arrow.Second}, lhsJSON) 548 defer lhs.Release() 549 rhs := c.getArr(&arrow.TimestampType{Unit: arrow.Second, TimeZone: "America/Phoenix"}, rhsJSON) 550 defer rhs.Release() 551 552 _, err := compute.CallFunction(c.ctx, op.fn, nil, &compute.ArrayDatum{lhs.Data()}, &compute.ArrayDatum{rhs.Data()}) 553 c.ErrorIs(err, arrow.ErrInvalid) 554 c.ErrorContains(err, "cannot compare timestamp with timezone to timestamp without timezone") 555 556 lhs = c.getArr(&arrow.TimestampType{Unit: arrow.Second, TimeZone: "America/New_York"}, lhsJSON) 557 defer lhs.Release() 558 rhs = c.getArr(&arrow.TimestampType{Unit: arrow.Second}, rhsJSON) 559 defer rhs.Release() 560 561 _, err = compute.CallFunction(c.ctx, op.fn, nil, &compute.ArrayDatum{lhs.Data()}, &compute.ArrayDatum{rhs.Data()}) 562 c.ErrorIs(err, arrow.ErrInvalid) 563 c.ErrorContains(err, "cannot compare timestamp with timezone to timestamp without timezone") 564 }) 565 }) 566 } 567 } 568 569 func (c *CompareTimestampSuite) TestScalarArray() { 570 const scalarStr = "1970-01-02" 571 const arrayJSON = `["1970-01-02", "2000-02-01", null, "1900-02-28"]` 572 573 checkArrCase := func(scType, arrayType arrow.DataType, op kernels.CompareOperator, expectedJSON, flipExpectedJSON string) { 574 scalarSide, err := scalar.MakeScalarParam(scalarStr, scType) 575 c.Require().NoError(err) 576 arraySide := c.getArr(arrayType, arrayJSON) 577 defer arraySide.Release() 578 579 expected := c.getArr(arrow.FixedWidthTypes.Boolean, expectedJSON) 580 defer expected.Release() 581 flipExpected := c.getArr(arrow.FixedWidthTypes.Boolean, flipExpectedJSON) 582 defer flipExpected.Release() 583 584 cases := []struct{ side1, side2, expected compute.Datum }{ 585 {compute.NewDatum(scalarSide), &compute.ArrayDatum{arraySide.Data()}, &compute.ArrayDatum{expected.Data()}}, 586 {&compute.ArrayDatum{arraySide.Data()}, compute.NewDatum(scalarSide), &compute.ArrayDatum{flipExpected.Data()}}, 587 } 588 589 for _, arrCase := range cases { 590 lhs, rhs := arrCase.side1, arrCase.side2 591 if arrow.TypeEqual(scType, arrayType) { 592 c.validateCompareDatum(op, lhs, rhs, arrCase.expected) 593 } else { 594 _, err := compute.CallFunction(c.ctx, op.String(), nil, lhs, rhs) 595 c.ErrorIs(err, arrow.ErrInvalid) 596 c.ErrorContains(err, "cannot compare timestamp with timezone to timestamp without timezone") 597 } 598 } 599 } 600 601 for _, unit := range arrow.TimeUnitValues { 602 c.Run(unit.String(), func() { 603 tests := []struct{ t0, t1 arrow.DataType }{ 604 {&arrow.TimestampType{Unit: unit}, &arrow.TimestampType{Unit: unit}}, 605 {&arrow.TimestampType{Unit: unit}, &arrow.TimestampType{Unit: unit, TimeZone: "utc"}}, 606 {&arrow.TimestampType{Unit: unit, TimeZone: "utc"}, &arrow.TimestampType{Unit: unit}}, 607 {&arrow.TimestampType{Unit: unit, TimeZone: "utc"}, &arrow.TimestampType{Unit: unit, TimeZone: "utc"}}, 608 } 609 for _, tt := range tests { 610 checkArrCase(tt.t0, tt.t1, kernels.CmpEQ, `[true, false, null, false]`, `[true, false, null, false]`) 611 checkArrCase(tt.t0, tt.t1, kernels.CmpNE, `[false, true, null, true]`, `[false, true, null, true]`) 612 checkArrCase(tt.t0, tt.t1, kernels.CmpLT, `[false, true, null, false]`, `[false, false, null, true]`) 613 checkArrCase(tt.t0, tt.t1, kernels.CmpLE, `[true, true, null, false]`, `[true, false, null, true]`) 614 checkArrCase(tt.t0, tt.t1, kernels.CmpGT, `[false, false, null, true]`, `[false, true, null, false]`) 615 checkArrCase(tt.t0, tt.t1, kernels.CmpGE, `[true, false, null, true]`, `[true, true, null, false]`) 616 } 617 }) 618 } 619 } 620 621 type CompareDecimalSuite struct { 622 CompareSuite 623 } 624 625 func (c *CompareDecimalSuite) TestArrayScalar() { 626 cases := []struct{ fn, exp string }{ 627 {"equal", `[true, false, false, null]`}, 628 {"not_equal", `[false, true, true, null]`}, 629 {"less", `[false, false, true, null]`}, 630 {"less_equal", `[true, false, true, null]`}, 631 {"greater", `[false, true, false, null]`}, 632 {"greater_equal", `[true, true, false, null]`}, 633 } 634 635 for _, id := range []arrow.Type{arrow.DECIMAL128, arrow.DECIMAL256} { 636 c.Run(id.String(), func() { 637 ty, _ := arrow.NewDecimalType(id, 3, 2) 638 639 lhsArr := c.getArr(ty, `["1.23", "2.34", "-1.23", null]`) 640 lhsFloatArr := c.getArr(arrow.PrimitiveTypes.Float64, `[1.23, 2.34, -1.23, null]`) 641 lhsIntLikeArr := c.getArr(ty, `["1.00", "2.00", "-1.00", null]`) 642 defer func() { 643 lhsArr.Release() 644 lhsFloatArr.Release() 645 lhsIntLikeArr.Release() 646 }() 647 648 lhs := &compute.ArrayDatum{lhsArr.Data()} 649 lhsFloat := &compute.ArrayDatum{lhsFloatArr.Data()} 650 lhsIntLike := &compute.ArrayDatum{lhsIntLikeArr.Data()} 651 652 rhs, _ := scalar.MakeScalarParam("1.23", ty) 653 rhsFloat := scalar.MakeScalar(float64(1.23)) 654 rhsInt := scalar.MakeScalar(int64(1)) 655 for _, tc := range cases { 656 c.Run(tc.fn, func() { 657 exp := c.getArr(arrow.FixedWidthTypes.Boolean, tc.exp) 658 defer exp.Release() 659 expected := &compute.ArrayDatum{exp.Data()} 660 661 checkScalarBinary(c.T(), tc.fn, lhs, compute.NewDatum(rhs), expected, nil) 662 checkScalarBinary(c.T(), tc.fn, lhsFloat, compute.NewDatum(rhs), expected, nil) 663 checkScalarBinary(c.T(), tc.fn, lhs, compute.NewDatum(rhsFloat), expected, nil) 664 checkScalarBinary(c.T(), tc.fn, lhsIntLike, compute.NewDatum(rhsInt), expected, nil) 665 }) 666 } 667 }) 668 } 669 } 670 671 func (c *CompareDecimalSuite) TestScalarArray() { 672 cases := []struct{ fn, exp string }{ 673 {"equal", `[true, false, false, null]`}, 674 {"not_equal", `[false, true, true, null]`}, 675 {"less", `[false, true, false, null]`}, 676 {"less_equal", `[true, true, false, null]`}, 677 {"greater", `[false, false, true, null]`}, 678 {"greater_equal", `[true, false, true, null]`}, 679 } 680 681 for _, id := range []arrow.Type{arrow.DECIMAL128, arrow.DECIMAL256} { 682 c.Run(id.String(), func() { 683 ty, _ := arrow.NewDecimalType(id, 3, 2) 684 685 rhsArr := c.getArr(ty, `["1.23", "2.34", "-1.23", null]`) 686 rhsFloatArr := c.getArr(arrow.PrimitiveTypes.Float64, `[1.23, 2.34, -1.23, null]`) 687 rhsIntLikeArr := c.getArr(ty, `["1.00", "2.00", "-1.00", null]`) 688 defer func() { 689 rhsArr.Release() 690 rhsFloatArr.Release() 691 rhsIntLikeArr.Release() 692 }() 693 694 rhs := &compute.ArrayDatum{rhsArr.Data()} 695 rhsFloat := &compute.ArrayDatum{rhsFloatArr.Data()} 696 rhsIntLike := &compute.ArrayDatum{rhsIntLikeArr.Data()} 697 698 lhs, _ := scalar.MakeScalarParam("1.23", ty) 699 lhsFloat := scalar.MakeScalar(float64(1.23)) 700 lhsInt := scalar.MakeScalar(int64(1)) 701 for _, tc := range cases { 702 c.Run(tc.fn, func() { 703 exp := c.getArr(arrow.FixedWidthTypes.Boolean, tc.exp) 704 defer exp.Release() 705 expected := &compute.ArrayDatum{exp.Data()} 706 707 checkScalarBinary(c.T(), tc.fn, compute.NewDatum(lhs), rhs, expected, nil) 708 checkScalarBinary(c.T(), tc.fn, compute.NewDatum(lhs), rhsFloat, expected, nil) 709 checkScalarBinary(c.T(), tc.fn, compute.NewDatum(lhsFloat), rhs, expected, nil) 710 checkScalarBinary(c.T(), tc.fn, compute.NewDatum(lhsInt), rhsIntLike, expected, nil) 711 }) 712 } 713 }) 714 } 715 } 716 717 func (c *CompareDecimalSuite) TestArrayArray() { 718 cases := []struct{ fn, exp string }{ 719 {"equal", `[true, false, false, true, false, false, null, null]`}, 720 {"not_equal", `[false, true, true, false, true, true, null, null]`}, 721 {"less", `[false, true, false, false, true, false, null, null]`}, 722 {"less_equal", `[true, true, false, true, true, false, null, null]`}, 723 {"greater", `[false, false, true, false, false, true, null, null]`}, 724 {"greater_equal", `[true, false, true, true, false, true, null, null]`}, 725 } 726 727 for _, id := range []arrow.Type{arrow.DECIMAL128, arrow.DECIMAL256} { 728 c.Run(id.String(), func() { 729 ty, _ := arrow.NewDecimalType(id, 3, 2) 730 731 lhsArr := c.getArr(ty, `["1.23", "1.23", "2.34", "-1.23", "-1.23", "1.23", "1.23", null]`) 732 lhsFloatArr := c.getArr(arrow.PrimitiveTypes.Float64, `[1.23, 1.23, 2.34, -1.23, -1.23, 1.23, 1.23, null]`) 733 lhsIntLikeArr := c.getArr(ty, `["1.00", "1.00", "2.00", "-1.00", "-1.00", "1.00", "1.00", null]`) 734 defer func() { 735 lhsArr.Release() 736 lhsFloatArr.Release() 737 lhsIntLikeArr.Release() 738 }() 739 740 lhs := &compute.ArrayDatum{lhsArr.Data()} 741 lhsFloat := &compute.ArrayDatum{lhsFloatArr.Data()} 742 lhsIntLike := &compute.ArrayDatum{lhsIntLikeArr.Data()} 743 744 rhsArr := c.getArr(ty, `["1.23", "2.34", "1.23", "-1.23", "1.23", "-1.23", null, "1.23"]`) 745 rhsFloatArr := c.getArr(arrow.PrimitiveTypes.Float64, `[1.23, 2.34, 1.23, -1.23, 1.23, -1.23, null, 1.23]`) 746 rhsIntArr := c.getArr(arrow.PrimitiveTypes.Int64, `[1, 2, 1, -1, 1, -1, null, 1]`) 747 defer func() { 748 rhsArr.Release() 749 rhsFloatArr.Release() 750 rhsIntArr.Release() 751 }() 752 753 rhs := &compute.ArrayDatum{rhsArr.Data()} 754 rhsFloat := &compute.ArrayDatum{rhsFloatArr.Data()} 755 rhsInt := &compute.ArrayDatum{rhsIntArr.Data()} 756 757 empty := c.getArr(ty, `[]`) 758 emptyExp := c.getArr(arrow.FixedWidthTypes.Boolean, `[]`) 759 null := c.getArr(ty, `[null]`) 760 nullExp := c.getArr(arrow.FixedWidthTypes.Boolean, `[null]`) 761 defer func() { 762 empty.Release() 763 emptyExp.Release() 764 null.Release() 765 nullExp.Release() 766 }() 767 768 for _, tc := range cases { 769 c.Run(tc.fn, func() { 770 exp := c.getArr(arrow.FixedWidthTypes.Boolean, tc.exp) 771 defer exp.Release() 772 expected := &compute.ArrayDatum{exp.Data()} 773 774 checkScalarBinary(c.T(), tc.fn, &compute.ArrayDatum{empty.Data()}, 775 &compute.ArrayDatum{empty.Data()}, &compute.ArrayDatum{emptyExp.Data()}, nil) 776 checkScalarBinary(c.T(), tc.fn, &compute.ArrayDatum{null.Data()}, 777 &compute.ArrayDatum{null.Data()}, &compute.ArrayDatum{nullExp.Data()}, nil) 778 checkScalarBinary(c.T(), tc.fn, lhs, rhs, expected, nil) 779 checkScalarBinary(c.T(), tc.fn, lhsFloat, rhs, expected, nil) 780 checkScalarBinary(c.T(), tc.fn, lhs, rhsFloat, expected, nil) 781 checkScalarBinary(c.T(), tc.fn, lhsIntLike, rhsInt, expected, nil) 782 }) 783 } 784 }) 785 } 786 } 787 788 func (c *CompareDecimalSuite) TestDiffParams() { 789 cases := []struct{ fn, exp string }{ 790 {"equal", `[true, false, false, true, false, false]`}, 791 {"not_equal", `[false, true, true, false, true, true]`}, 792 {"less", `[false, true, false, false, true, false]`}, 793 {"less_equal", `[true, true, false, true, true, false]`}, 794 {"greater", `[false, false, true, false, false, true]`}, 795 {"greater_equal", `[true, false, true, true, false, true]`}, 796 } 797 798 for _, id := range []arrow.Type{arrow.DECIMAL128, arrow.DECIMAL256} { 799 c.Run(id.String(), func() { 800 ty1, _ := arrow.NewDecimalType(id, 3, 2) 801 ty2, _ := arrow.NewDecimalType(id, 4, 3) 802 803 lhsArr := c.getArr(ty1, `["1.23", "1.23", "2.34", "-1.23", "-1.23", "1.23"]`) 804 rhsArr := c.getArr(ty2, `["1.230", "2.340", "1.230", "-1.230", "1.230", "-1.230"]`) 805 defer func() { 806 lhsArr.Release() 807 rhsArr.Release() 808 }() 809 810 lhs := &compute.ArrayDatum{lhsArr.Data()} 811 rhs := &compute.ArrayDatum{rhsArr.Data()} 812 813 for _, tc := range cases { 814 c.Run(tc.fn, func() { 815 exp := c.getArr(arrow.FixedWidthTypes.Boolean, tc.exp) 816 defer exp.Release() 817 expected := &compute.ArrayDatum{exp.Data()} 818 819 checkScalarBinary(c.T(), tc.fn, lhs, rhs, expected, nil) 820 }) 821 } 822 }) 823 } 824 } 825 826 type CompareFixedSizeBinary struct { 827 CompareSuite 828 } 829 830 type fsbCompareCase struct { 831 lhsType, rhsType arrow.DataType 832 lhs, rhs string 833 // index into cases[...].exp 834 resultIdx int 835 } 836 837 func (c *CompareFixedSizeBinary) TestArrayScalar() { 838 ty1 := &arrow.FixedSizeBinaryType{ByteWidth: 3} 839 ty2 := &arrow.FixedSizeBinaryType{ByteWidth: 1} 840 841 cases := []struct { 842 fn string 843 exp []string 844 }{ 845 {"equal", []string{ 846 `[false, true, false, null]`, 847 `[false, false, false, null]`, 848 `[false, false, false, null]`}}, 849 {"not_equal", []string{ 850 `[true, false, true, null]`, 851 `[true, true, true, null]`, 852 `[true, true, true, null]`}}, 853 {"less", []string{ 854 `[true, false, false, null]`, 855 `[true, true, true, null]`, 856 `[true, false, false, null]`}}, 857 {"less_equal", []string{ 858 `[true, true, false, null]`, 859 `[true, true, true, null]`, 860 `[true, false, false, null]`}}, 861 {"greater", []string{ 862 `[false, false, true, null]`, 863 `[false, false, false, null]`, 864 `[false, true, true, null]`}}, 865 {"greater_equal", []string{ 866 `[false, true, true, null]`, 867 `[false, false, false, null]`, 868 `[false, true, true, null]`}}, 869 } 870 871 // base64 encoding 872 const ( 873 valAba = `YWJh` 874 valAbc = `YWJj` 875 valAbd = `YWJk` 876 valA = `YQ==` 877 valB = `Yg==` 878 valC = `Yw==` 879 ) 880 881 const ( 882 lhs1bin = `["` + valAba + `","` + valAbc + `","` + valAbd + `", null]` 883 lhs1 = `["aba", "abc", "abd", null]` 884 rhs1 = "abc" 885 lhs2bin = `["` + valA + `","` + valB + `","` + valC + `", null]` 886 rhs2 = "b" 887 ) 888 889 types := []fsbCompareCase{ 890 {ty1, ty1, lhs1bin, rhs1, 0}, 891 {ty2, ty2, lhs2bin, rhs2, 0}, 892 {ty1, ty2, lhs1bin, rhs2, 1}, 893 {ty2, ty1, lhs2bin, rhs1, 2}, 894 {ty1, arrow.BinaryTypes.Binary, lhs1bin, rhs1, 0}, 895 {arrow.BinaryTypes.Binary, ty1, lhs1bin, rhs1, 0}, 896 {ty1, arrow.BinaryTypes.LargeBinary, lhs1bin, rhs1, 0}, 897 {arrow.BinaryTypes.LargeBinary, ty1, lhs1bin, rhs1, 0}, 898 {ty1, arrow.BinaryTypes.String, lhs1bin, rhs1, 0}, 899 {arrow.BinaryTypes.String, ty1, lhs1, rhs1, 0}, 900 {ty1, arrow.BinaryTypes.LargeString, lhs1bin, rhs1, 0}, 901 {arrow.BinaryTypes.LargeString, ty1, lhs1, rhs1, 0}, 902 } 903 904 expNull := c.getArr(arrow.FixedWidthTypes.Boolean, `[null]`) 905 defer expNull.Release() 906 907 for _, op := range cases { 908 c.Run(op.fn, func() { 909 for _, tc := range types { 910 lhs := c.getArr(tc.lhsType, tc.lhs) 911 defer lhs.Release() 912 rhs, _ := scalar.MakeScalarParam(tc.rhs, tc.rhsType) 913 exp := c.getArr(arrow.FixedWidthTypes.Boolean, op.exp[tc.resultIdx]) 914 defer exp.Release() 915 916 expected := &compute.ArrayDatum{exp.Data()} 917 918 null := c.getArr(tc.lhsType, `[null]`) 919 defer null.Release() 920 scNull := scalar.MakeNullScalar(tc.rhsType) 921 922 checkScalarBinary(c.T(), op.fn, &compute.ArrayDatum{null.Data()}, compute.NewDatum(scNull), 923 &compute.ArrayDatum{expNull.Data()}, nil) 924 checkScalarBinary(c.T(), op.fn, &compute.ArrayDatum{lhs.Data()}, 925 compute.NewDatum(rhs), expected, nil) 926 } 927 }) 928 } 929 } 930 931 func (c *CompareFixedSizeBinary) TestScalarArray() { 932 ty1 := &arrow.FixedSizeBinaryType{ByteWidth: 3} 933 ty2 := &arrow.FixedSizeBinaryType{ByteWidth: 1} 934 935 cases := []struct { 936 fn string 937 exp []string 938 }{ 939 {"equal", []string{ 940 `[false, true, false, null]`, 941 `[false, false, false, null]`, 942 `[false, false, false, null]`}}, 943 {"not_equal", []string{ 944 `[true, false, true, null]`, 945 `[true, true, true, null]`, 946 `[true, true, true, null]`}}, 947 {"less", []string{ 948 `[false, false, true, null]`, 949 `[false, true, true, null]`, 950 `[false, false, false, null]`}}, 951 {"less_equal", []string{ 952 `[false, true, true, null]`, 953 `[false, true, true, null]`, 954 `[false, false, false, null]`}}, 955 {"greater", []string{ 956 `[true, false, false, null]`, 957 `[true, false, false, null]`, 958 `[true, true, true, null]`}}, 959 {"greater_equal", []string{ 960 `[true, true, false, null]`, 961 `[true, false, false, null]`, 962 `[true, true, true, null]`}}, 963 } 964 965 // base64 encoding 966 const ( 967 valAba = `YWJh` 968 valAbc = `YWJj` 969 valAbd = `YWJk` 970 valA = `YQ==` 971 valB = `Yg==` 972 valC = `Yw==` 973 ) 974 975 const ( 976 lhs1 = "abc" 977 rhs1bin = `["` + valAba + `","` + valAbc + `","` + valAbd + `", null]` 978 rhs1 = `["aba", "abc", "abd", null]` 979 lhs2 = "b" 980 rhs2bin = `["` + valA + `","` + valB + `","` + valC + `", null]` 981 rhs2 = `["a", "b", "c", null]` 982 ) 983 984 types := []fsbCompareCase{ 985 {ty1, ty1, lhs1, rhs1bin, 0}, 986 {ty2, ty2, lhs2, rhs2bin, 0}, 987 {ty1, ty2, lhs1, rhs2bin, 1}, 988 {ty2, ty1, lhs2, rhs1bin, 2}, 989 {ty1, arrow.BinaryTypes.Binary, lhs1, rhs1bin, 0}, 990 {arrow.BinaryTypes.Binary, ty1, lhs1, rhs1bin, 0}, 991 {ty1, arrow.BinaryTypes.LargeBinary, lhs1, rhs1bin, 0}, 992 {arrow.BinaryTypes.LargeBinary, ty1, lhs1, rhs1bin, 0}, 993 {ty1, arrow.BinaryTypes.String, lhs1, rhs1, 0}, 994 {arrow.BinaryTypes.String, ty1, lhs1, rhs1bin, 0}, 995 {ty1, arrow.BinaryTypes.LargeString, lhs1, rhs1, 0}, 996 {arrow.BinaryTypes.LargeString, ty1, lhs1, rhs1bin, 0}, 997 } 998 999 expNull := c.getArr(arrow.FixedWidthTypes.Boolean, `[null]`) 1000 defer expNull.Release() 1001 1002 for _, op := range cases { 1003 c.Run(op.fn, func() { 1004 for _, tc := range types { 1005 lhs, _ := scalar.MakeScalarParam(tc.lhs, tc.lhsType) 1006 rhs := c.getArr(tc.rhsType, tc.rhs) 1007 defer rhs.Release() 1008 exp := c.getArr(arrow.FixedWidthTypes.Boolean, op.exp[tc.resultIdx]) 1009 defer exp.Release() 1010 1011 expected := &compute.ArrayDatum{exp.Data()} 1012 1013 null := c.getArr(tc.rhsType, `[null]`) 1014 defer null.Release() 1015 scNull := scalar.MakeNullScalar(tc.lhsType) 1016 1017 checkScalarBinary(c.T(), op.fn, compute.NewDatum(scNull), &compute.ArrayDatum{null.Data()}, 1018 &compute.ArrayDatum{expNull.Data()}, nil) 1019 checkScalarBinary(c.T(), op.fn, compute.NewDatum(lhs), 1020 &compute.ArrayDatum{rhs.Data()}, expected, nil) 1021 } 1022 }) 1023 } 1024 } 1025 1026 func (c *CompareFixedSizeBinary) TestArrayArray() { 1027 ty1 := &arrow.FixedSizeBinaryType{ByteWidth: 3} 1028 ty2 := &arrow.FixedSizeBinaryType{ByteWidth: 1} 1029 1030 cases := []struct { 1031 fn string 1032 exp []string 1033 }{ 1034 {"equal", []string{ 1035 `[true, false, false, null, null]`, 1036 `[true, false, false, null, null]`, 1037 `[true, false, false, null, null]`, 1038 `[true, false, false, null, null]`, 1039 `[false, false, false, null, null]`, 1040 `[false, false, false, null, null]`}}, 1041 {"not_equal", []string{ 1042 `[false, true, true, null, null]`, 1043 `[false, true, true, null, null]`, 1044 `[false, true, true, null, null]`, 1045 `[false, true, true, null, null]`, 1046 `[true, true, true, null, null]`, 1047 `[true, true, true, null, null]`}}, 1048 {"less", []string{ 1049 `[false, true, false, null, null]`, 1050 `[false, false, true, null, null]`, 1051 `[false, true, false, null, null]`, 1052 `[false, false, true, null, null]`, 1053 `[false, true, true, null, null]`, 1054 `[true, true, false, null, null]`}}, 1055 {"less_equal", []string{ 1056 `[true, true, false, null, null]`, 1057 `[true, false, true, null, null]`, 1058 `[true, true, false, null, null]`, 1059 `[true, false, true, null, null]`, 1060 `[false, true, true, null, null]`, 1061 `[true, true, false, null, null]`}}, 1062 {"greater", []string{ 1063 `[false, false, true, null, null]`, 1064 `[false, true, false, null, null]`, 1065 `[false, false, true, null, null]`, 1066 `[false, true, false, null, null]`, 1067 `[true, false, false, null, null]`, 1068 `[false, false, true, null, null]`}}, 1069 {"greater_equal", []string{ 1070 `[true, false, true, null, null]`, 1071 `[true, true, false, null, null]`, 1072 `[true, false, true, null, null]`, 1073 `[true, true, false, null, null]`, 1074 `[true, false, false, null, null]`, 1075 `[false, false, true, null, null]`}}, 1076 } 1077 1078 // base64 encoding 1079 const ( 1080 valAbc = `YWJj` 1081 valAbd = `YWJk` 1082 valA = `YQ==` 1083 valC = `Yw==` 1084 valD = `ZA==` 1085 ) 1086 1087 const ( 1088 lhs1bin = `["` + valAbc + `","` + valAbc + `","` + valAbd + `", null, "` + valAbc + `"]` 1089 rhs1bin = `["` + valAbc + `","` + valAbd + `","` + valAbc + `","` + valAbc + `", null]` 1090 lhs1 = `["abc", "abc", "abd", null, "abc"]` 1091 rhs1 = `["abc", "abd", "abc", "abc", null]` 1092 lhs2bin = `["` + valA + `","` + valA + `","` + valD + `", null, "` + valA + `"]` 1093 rhs2bin = `["` + valA + `","` + valD + `","` + valC + `","` + valA + `", null]` 1094 ) 1095 1096 types := []fsbCompareCase{ 1097 {ty1, ty1, lhs1bin, rhs1bin, 0}, 1098 {ty1, ty1, rhs1bin, lhs1bin, 1}, 1099 {ty2, ty2, lhs2bin, rhs2bin, 2}, 1100 {ty2, ty2, rhs2bin, lhs2bin, 3}, 1101 {ty1, ty2, lhs1bin, rhs2bin, 4}, 1102 {ty2, ty1, lhs2bin, rhs1bin, 5}, 1103 {ty1, arrow.BinaryTypes.Binary, lhs1bin, rhs1bin, 0}, 1104 {arrow.BinaryTypes.Binary, ty1, lhs1bin, rhs1bin, 0}, 1105 {ty1, arrow.BinaryTypes.LargeBinary, lhs1bin, rhs1bin, 0}, 1106 {arrow.BinaryTypes.LargeBinary, ty1, lhs1bin, rhs1bin, 0}, 1107 {ty1, arrow.BinaryTypes.String, lhs1bin, rhs1, 0}, 1108 {arrow.BinaryTypes.String, ty1, lhs1, rhs1bin, 0}, 1109 {ty1, arrow.BinaryTypes.LargeString, lhs1bin, rhs1, 0}, 1110 {arrow.BinaryTypes.LargeString, ty1, lhs1, rhs1bin, 0}, 1111 } 1112 1113 expEmpty := c.getArr(arrow.FixedWidthTypes.Boolean, `[]`) 1114 defer expEmpty.Release() 1115 expNull := c.getArr(arrow.FixedWidthTypes.Boolean, `[null]`) 1116 defer expNull.Release() 1117 1118 for _, op := range cases { 1119 c.Run(op.fn, func() { 1120 for _, tc := range types { 1121 lhs := c.getArr(tc.lhsType, tc.lhs) 1122 defer lhs.Release() 1123 rhs := c.getArr(tc.rhsType, tc.rhs) 1124 defer rhs.Release() 1125 exp := c.getArr(arrow.FixedWidthTypes.Boolean, op.exp[tc.resultIdx]) 1126 defer exp.Release() 1127 1128 expected := &compute.ArrayDatum{exp.Data()} 1129 1130 lhsEmpty := c.getArr(tc.lhsType, `[]`) 1131 defer lhsEmpty.Release() 1132 rhsEmpty := c.getArr(tc.rhsType, `[]`) 1133 defer rhsEmpty.Release() 1134 lhsNull := c.getArr(tc.lhsType, `[null]`) 1135 defer lhsNull.Release() 1136 rhsNull := c.getArr(tc.rhsType, `[null]`) 1137 defer rhsNull.Release() 1138 1139 checkScalarBinary(c.T(), op.fn, &compute.ArrayDatum{lhsEmpty.Data()}, &compute.ArrayDatum{rhsEmpty.Data()}, 1140 &compute.ArrayDatum{expEmpty.Data()}, nil) 1141 checkScalarBinary(c.T(), op.fn, &compute.ArrayDatum{lhsNull.Data()}, &compute.ArrayDatum{rhsNull.Data()}, 1142 &compute.ArrayDatum{expNull.Data()}, nil) 1143 checkScalarBinary(c.T(), op.fn, &compute.ArrayDatum{lhs.Data()}, 1144 &compute.ArrayDatum{rhs.Data()}, expected, nil) 1145 } 1146 }) 1147 } 1148 } 1149 1150 type CompareStringSuite struct { 1151 CompareSuite 1152 } 1153 1154 func (c *CompareStringSuite) TestSimpleCompareArrayScalar() { 1155 one := compute.NewDatum(scalar.MakeScalar("one")) 1156 1157 dt := arrow.BinaryTypes.String 1158 1159 op := kernels.CmpEQ 1160 c.validateCompareArrScalar(op, dt, `[]`, one, `[]`) 1161 c.validateCompareArrScalar(op, dt, `[null]`, one, `[null]`) 1162 c.validateCompareArrScalar(op, dt, `["zero", "zero", "one", "one", "two", "two"]`, one, 1163 `[false, false, true, true, false, false]`) 1164 c.validateCompareArrScalar(op, dt, `["zero", "one", "two", "three", "four", "five"]`, one, 1165 `[false, true, false, false, false, false]`) 1166 c.validateCompareArrScalar(op, dt, `["five", "four", "three", "two", "one", "zero"]`, one, 1167 `[false, false, false, false, true, false]`) 1168 c.validateCompareArrScalar(op, dt, `[null, "zero", "one", "one"]`, one, `[null, false, true, true]`) 1169 1170 na := compute.NewDatum(scalar.MakeNullScalar(dt)) 1171 c.validateCompareArrScalar(op, dt, `[null, "zero", "one", "one"]`, na, `[null, null, null, null]`) 1172 c.validateCompareScalarArr(op, dt, na, `[null, "zero", "one", "one"]`, `[null, null, null, null]`) 1173 1174 op = kernels.CmpNE 1175 c.validateCompareArrScalar(op, dt, `[]`, one, `[]`) 1176 c.validateCompareArrScalar(op, dt, `[null]`, one, `[null]`) 1177 c.validateCompareArrScalar(op, dt, `["zero", "zero", "one", "one", "two", "two"]`, one, 1178 `[true, true, false, false, true, true]`) 1179 c.validateCompareArrScalar(op, dt, `["zero", "one", "two", "three", "four", "five"]`, one, 1180 `[true, false, true, true, true, true]`) 1181 c.validateCompareArrScalar(op, dt, `["five", "four", "three", "two", "one", "zero"]`, one, 1182 `[true, true, true, true, false, true]`) 1183 c.validateCompareArrScalar(op, dt, `[null, "zero", "one", "one"]`, one, `[null, true, false, false]`) 1184 } 1185 1186 func (c *CompareStringSuite) validateCompareComputed(op kernels.CompareOperator, lhs, rhs compute.Datum) { 1187 var expected compute.Datum 1188 1189 hasScalar := lhs.Kind() == compute.KindScalar || rhs.Kind() == compute.KindScalar 1190 if hasScalar { 1191 expected = simpleScalarArrayCompareString(c.mem, op, lhs, rhs) 1192 } else { 1193 expected = simpleArrArrCompare[string](c.mem, op, lhs, rhs) 1194 } 1195 1196 defer expected.Release() 1197 c.CompareSuite.validateCompareDatum(op, lhs, rhs, expected) 1198 } 1199 1200 func (c *CompareStringSuite) TestRandomCompareArrayArray() { 1201 rng := gen.NewRandomArrayGenerator(0x5416447, c.mem) 1202 for i := 3; i < 5; i++ { 1203 c.Run(fmt.Sprintf("len=%d", 1<<i), func() { 1204 for _, nullProb := range []float64{0.0, 0.01, 0.1, 0.25, 0.5, 1.0} { 1205 c.Run(fmt.Sprintf("nullprob=%0.2f", nullProb), func() { 1206 for _, op := range []kernels.CompareOperator{kernels.CmpEQ, kernels.CmpNE} { 1207 c.Run(op.String(), func() { 1208 length := int64(1 << i) 1209 lhs := rng.String(length<<i, 0, 16, nullProb) 1210 defer lhs.Release() 1211 rhs := rng.String(length<<i, 0, 16, nullProb) 1212 defer rhs.Release() 1213 1214 c.validateCompareComputed(op, 1215 &compute.ArrayDatum{lhs.Data()}, 1216 &compute.ArrayDatum{rhs.Data()}) 1217 }) 1218 } 1219 }) 1220 } 1221 }) 1222 } 1223 } 1224 1225 func TestComparisons(t *testing.T) { 1226 suite.Run(t, new(NumericCompareSuite[int8])) 1227 suite.Run(t, new(NumericCompareSuite[int16])) 1228 suite.Run(t, new(NumericCompareSuite[int32])) 1229 suite.Run(t, new(NumericCompareSuite[int64])) 1230 suite.Run(t, new(NumericCompareSuite[uint8])) 1231 suite.Run(t, new(NumericCompareSuite[uint16])) 1232 suite.Run(t, new(NumericCompareSuite[uint32])) 1233 suite.Run(t, new(NumericCompareSuite[uint64])) 1234 suite.Run(t, new(NumericCompareSuite[float32])) 1235 suite.Run(t, new(NumericCompareSuite[float64])) 1236 suite.Run(t, new(CompareTimestampSuite)) 1237 suite.Run(t, new(CompareDecimalSuite)) 1238 suite.Run(t, new(CompareFixedSizeBinary)) 1239 suite.Run(t, new(CompareStringSuite)) 1240 } 1241 1242 func TestCompareKernelsDispatchBest(t *testing.T) { 1243 tests := []struct { 1244 origLeft, origRight arrow.DataType 1245 expectLeft, expectRight arrow.DataType 1246 }{ 1247 {arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int32}, 1248 {arrow.PrimitiveTypes.Int32, arrow.Null, arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int32}, 1249 {arrow.Null, arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int32}, 1250 1251 {arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int8, arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int32}, 1252 {arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int16, arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int32}, 1253 {arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int64, arrow.PrimitiveTypes.Int64, arrow.PrimitiveTypes.Int64}, 1254 1255 {arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Uint8, arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int32}, 1256 {arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Uint16, arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int32}, 1257 {arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Uint32, arrow.PrimitiveTypes.Int64, arrow.PrimitiveTypes.Int64}, 1258 {arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Uint64, arrow.PrimitiveTypes.Int64, arrow.PrimitiveTypes.Int64}, 1259 1260 {arrow.PrimitiveTypes.Uint8, arrow.PrimitiveTypes.Uint8, arrow.PrimitiveTypes.Uint8, arrow.PrimitiveTypes.Uint8}, 1261 {arrow.PrimitiveTypes.Uint8, arrow.PrimitiveTypes.Uint16, arrow.PrimitiveTypes.Uint16, arrow.PrimitiveTypes.Uint16}, 1262 1263 {arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Float32, arrow.PrimitiveTypes.Float32, arrow.PrimitiveTypes.Float32}, 1264 {arrow.PrimitiveTypes.Float32, arrow.PrimitiveTypes.Int64, arrow.PrimitiveTypes.Float32, arrow.PrimitiveTypes.Float32}, 1265 {arrow.PrimitiveTypes.Float64, arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Float64, arrow.PrimitiveTypes.Float64}, 1266 1267 {&arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int8, ValueType: arrow.PrimitiveTypes.Float64}, arrow.PrimitiveTypes.Float64, 1268 arrow.PrimitiveTypes.Float64, arrow.PrimitiveTypes.Float64}, 1269 {&arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int8, ValueType: arrow.PrimitiveTypes.Float64}, arrow.PrimitiveTypes.Int16, 1270 arrow.PrimitiveTypes.Float64, arrow.PrimitiveTypes.Float64}, 1271 1272 {arrow.FixedWidthTypes.Timestamp_us, arrow.FixedWidthTypes.Date64, arrow.FixedWidthTypes.Timestamp_us, arrow.FixedWidthTypes.Timestamp_us}, 1273 {arrow.FixedWidthTypes.Timestamp_ms, arrow.FixedWidthTypes.Timestamp_us, arrow.FixedWidthTypes.Timestamp_us, arrow.FixedWidthTypes.Timestamp_us}, 1274 1275 {arrow.BinaryTypes.String, arrow.BinaryTypes.Binary, arrow.BinaryTypes.Binary, arrow.BinaryTypes.Binary}, 1276 {arrow.BinaryTypes.LargeString, arrow.BinaryTypes.Binary, arrow.BinaryTypes.LargeBinary, arrow.BinaryTypes.LargeBinary}, 1277 {arrow.BinaryTypes.LargeString, &arrow.FixedSizeBinaryType{ByteWidth: 2}, arrow.BinaryTypes.LargeBinary, arrow.BinaryTypes.LargeBinary}, 1278 {arrow.BinaryTypes.Binary, &arrow.FixedSizeBinaryType{ByteWidth: 2}, arrow.BinaryTypes.Binary, arrow.BinaryTypes.Binary}, 1279 {&arrow.FixedSizeBinaryType{ByteWidth: 4}, &arrow.FixedSizeBinaryType{ByteWidth: 2}, 1280 &arrow.FixedSizeBinaryType{ByteWidth: 4}, &arrow.FixedSizeBinaryType{ByteWidth: 2}}, 1281 1282 {&arrow.Decimal128Type{Precision: 3, Scale: 2}, &arrow.Decimal128Type{Precision: 6, Scale: 3}, 1283 &arrow.Decimal128Type{Precision: 4, Scale: 3}, &arrow.Decimal128Type{Precision: 6, Scale: 3}}, 1284 {&arrow.Decimal128Type{Precision: 3, Scale: 2}, &arrow.Decimal256Type{Precision: 3, Scale: 2}, 1285 &arrow.Decimal256Type{Precision: 3, Scale: 2}, &arrow.Decimal256Type{Precision: 3, Scale: 2}}, 1286 {&arrow.Decimal128Type{Precision: 3, Scale: 2}, arrow.PrimitiveTypes.Float64, arrow.PrimitiveTypes.Float64, arrow.PrimitiveTypes.Float64}, 1287 {arrow.PrimitiveTypes.Float64, &arrow.Decimal128Type{Precision: 3, Scale: 2}, arrow.PrimitiveTypes.Float64, arrow.PrimitiveTypes.Float64}, 1288 {&arrow.Decimal128Type{Precision: 3, Scale: 2}, arrow.PrimitiveTypes.Int64, 1289 &arrow.Decimal128Type{Precision: 3, Scale: 2}, &arrow.Decimal128Type{Precision: 21, Scale: 2}}, 1290 {arrow.PrimitiveTypes.Int64, &arrow.Decimal128Type{Precision: 3, Scale: 2}, 1291 &arrow.Decimal128Type{Precision: 21, Scale: 2}, &arrow.Decimal128Type{Precision: 3, Scale: 2}}, 1292 } 1293 1294 for _, name := range []string{"equal", "not_equal", "less", "less_equal", "greater", "greater_equal"} { 1295 t.Run(name, func(t *testing.T) { 1296 for _, tt := range tests { 1297 CheckDispatchBest(t, name, []arrow.DataType{tt.origLeft, tt.origRight}, 1298 []arrow.DataType{tt.expectLeft, tt.expectRight}) 1299 } 1300 }) 1301 } 1302 } 1303 1304 func TestCompareGreaterWithImplicitCasts(t *testing.T) { 1305 mem := memory.NewCheckedAllocator(memory.DefaultAllocator) 1306 defer mem.AssertSize(t, 0) 1307 1308 getArr := func(ty arrow.DataType, str string) arrow.Array { 1309 arr, _, err := array.FromJSON(mem, ty, strings.NewReader(str), array.WithUseNumber()) 1310 require.NoError(t, err) 1311 return arr 1312 } 1313 1314 check := func(ty1 arrow.DataType, str1 string, ty2 arrow.DataType, str2 string, exp string) { 1315 arr1, arr2 := getArr(ty1, str1), getArr(ty2, str2) 1316 arrExp := getArr(arrow.FixedWidthTypes.Boolean, exp) 1317 1318 checkScalarBinary(t, "greater", compute.NewDatumWithoutOwning(arr1), 1319 compute.NewDatumWithoutOwning(arr2), 1320 compute.NewDatumWithoutOwning(arrExp), nil) 1321 1322 arr1.Release() 1323 arr2.Release() 1324 arrExp.Release() 1325 } 1326 1327 tests := []struct { 1328 ty1, ty2 arrow.DataType 1329 str1, str2 string 1330 exp string 1331 }{ 1332 {arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Float64, 1333 `[0, 1, 2, null]`, `[0.5, 1.0, 1.5, 2.0]`, `[false, false, true, null]`}, 1334 {arrow.PrimitiveTypes.Int8, arrow.PrimitiveTypes.Uint32, 1335 `[-16, 0, 16, null]`, `[3, 4, 5, 7]`, `[false, false, true, null]`}, 1336 {arrow.PrimitiveTypes.Int8, arrow.PrimitiveTypes.Uint8, 1337 `[-16, 0, 16, null]`, `[255, 254, 1, 0]`, `[false, false, true, null]`}, 1338 {&arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.PrimitiveTypes.Int32}, 1339 arrow.PrimitiveTypes.Uint32, `[0, 1, 2, null]`, `[3, 4, 5, 7]`, `[false, false, false, null]`}, 1340 {&arrow.TimestampType{Unit: arrow.Second}, arrow.FixedWidthTypes.Date64, 1341 `["1970-01-01", "2000-02-29", "1900-02-28"]`, `[86400000, 0, 86400000]`, 1342 `[false, true, false]`}, 1343 {&arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: arrow.PrimitiveTypes.Int8}, 1344 arrow.PrimitiveTypes.Uint32, `[3, -3, -28, null]`, `[3, 4, 5, 7]`, 1345 `[false, false, false, null]`}, 1346 } 1347 1348 for _, tt := range tests { 1349 check(tt.ty1, tt.str1, tt.ty2, tt.str2, tt.exp) 1350 } 1351 } 1352 1353 func TestCompareGreaterWithImplicitCastUint64EdgeCase(t *testing.T) { 1354 // int64 is as wide as we can promote 1355 CheckDispatchBest(t, "greater", 1356 []arrow.DataType{arrow.PrimitiveTypes.Int8, arrow.PrimitiveTypes.Uint64}, 1357 []arrow.DataType{arrow.PrimitiveTypes.Int64, arrow.PrimitiveTypes.Int64}) 1358 1359 mem := memory.NewCheckedAllocator(memory.DefaultAllocator) 1360 defer mem.AssertSize(t, 0) 1361 1362 getArr := func(ty arrow.DataType, str string) arrow.Array { 1363 arr, _, err := array.FromJSON(mem, ty, strings.NewReader(str), array.WithUseNumber()) 1364 require.NoError(t, err) 1365 return arr 1366 } 1367 1368 // this works sometimes 1369 neg := getArr(arrow.PrimitiveTypes.Int8, `[-1]`) 1370 defer neg.Release() 1371 zero := getArr(arrow.PrimitiveTypes.Uint64, `[0]`) 1372 defer zero.Release() 1373 res := getArr(arrow.FixedWidthTypes.Boolean, `[false]`) 1374 defer res.Release() 1375 1376 checkScalarBinary(t, "greater", compute.NewDatumWithoutOwning(neg), 1377 compute.NewDatumWithoutOwning(zero), compute.NewDatumWithoutOwning(res), nil) 1378 1379 // ... but it can result in impossible implicit casts in the presence of uint64 1380 // since some uint64 values cannot be cast to int64 1381 neg = getArr(arrow.PrimitiveTypes.Int64, `[-1]`) 1382 defer neg.Release() 1383 big := getArr(arrow.PrimitiveTypes.Uint64, `[18446744073709551615]`) 1384 defer big.Release() 1385 1386 _, err := compute.CallFunction(context.TODO(), "greater", nil, compute.NewDatumWithoutOwning(neg), compute.NewDatumWithoutOwning(big)) 1387 assert.ErrorIs(t, err, arrow.ErrInvalid) 1388 } 1389 1390 const benchSeed = 0x94378165 1391 1392 func benchArrayScalar(b *testing.B, sz int, nullprob float64, op string, dt arrow.DataType) { 1393 b.Run(dt.String(), func(b *testing.B) { 1394 rng := gen.NewRandomArrayGenerator(benchSeed, memory.DefaultAllocator) 1395 arr := rng.ArrayOf(dt.ID(), int64(sz), nullprob) 1396 defer arr.Release() 1397 s := rng.ArrayOf(dt.ID(), 1, 0) 1398 defer s.Release() 1399 sc, _ := scalar.GetScalar(s, 0) 1400 1401 lhs := compute.NewDatumWithoutOwning(arr) 1402 rhs := compute.NewDatumWithoutOwning(sc) 1403 1404 var nbytes int64 1405 switch dt.ID() { 1406 case arrow.STRING: 1407 nbytes = int64(len(arr.(*array.String).ValueBytes()) + sc.(*scalar.String).Value.Len()) 1408 default: 1409 nbytes = int64(arr.Data().Buffers()[1].Len() + len(sc.(scalar.PrimitiveScalar).Data())) 1410 } 1411 ctx := context.Background() 1412 b.ResetTimer() 1413 b.SetBytes(nbytes) 1414 for n := 0; n < b.N; n++ { 1415 result, err := compute.CallFunction(ctx, op, nil, lhs, rhs) 1416 if err != nil { 1417 b.Fatal(err) 1418 } 1419 result.Release() 1420 } 1421 }) 1422 } 1423 1424 func benchArrayArray(b *testing.B, sz int, nullprob float64, op string, dt arrow.DataType) { 1425 b.Run(dt.String(), func(b *testing.B) { 1426 rng := gen.NewRandomArrayGenerator(benchSeed, memory.DefaultAllocator) 1427 lhsArr := rng.ArrayOf(dt.ID(), int64(sz), nullprob) 1428 defer lhsArr.Release() 1429 rhsArr := rng.ArrayOf(dt.ID(), int64(sz), nullprob) 1430 defer rhsArr.Release() 1431 1432 lhs, rhs := compute.NewDatumWithoutOwning(lhsArr), compute.NewDatumWithoutOwning(rhsArr) 1433 var nbytes int64 1434 switch dt.ID() { 1435 case arrow.STRING: 1436 nbytes = int64(len(lhsArr.(*array.String).ValueBytes()) + len(rhsArr.(*array.String).ValueBytes())) 1437 default: 1438 nbytes = int64(lhsArr.Data().Buffers()[1].Len() + rhsArr.Data().Buffers()[1].Len()) 1439 } 1440 ctx := context.Background() 1441 b.ResetTimer() 1442 b.SetBytes(nbytes) 1443 for n := 0; n < b.N; n++ { 1444 result, err := compute.CallFunction(ctx, op, nil, lhs, rhs) 1445 if err != nil { 1446 b.Fatal(err) 1447 } 1448 result.Release() 1449 } 1450 }) 1451 } 1452 1453 func BenchmarkCompare(b *testing.B) { 1454 var ( 1455 sizes = []int{CpuCacheSizes[0]} 1456 nullProb = []float64{0.0001, 0.01, 0.1, 0.5, 1, 0} 1457 ) 1458 1459 b.Run("GreaterArrayScalar", func(b *testing.B) { 1460 for _, sz := range sizes { 1461 b.Run(fmt.Sprintf("size=%d", sz), func(b *testing.B) { 1462 for _, np := range nullProb { 1463 b.Run(fmt.Sprintf("nullprob=%f", np), func(b *testing.B) { 1464 benchArrayScalar(b, sz, np, kernels.CmpGT.String(), arrow.PrimitiveTypes.Int64) 1465 benchArrayScalar(b, sz, np, kernels.CmpGT.String(), arrow.BinaryTypes.String) 1466 }) 1467 } 1468 }) 1469 } 1470 }) 1471 1472 b.Run("GreaterArrayArray", func(b *testing.B) { 1473 for _, sz := range sizes { 1474 b.Run(fmt.Sprintf("size=%d", sz), func(b *testing.B) { 1475 for _, np := range nullProb { 1476 b.Run(fmt.Sprintf("nullprob=%f", np), func(b *testing.B) { 1477 benchArrayArray(b, sz, np, kernels.CmpGT.String(), arrow.PrimitiveTypes.Int64) 1478 benchArrayArray(b, sz, np, kernels.CmpGT.String(), arrow.BinaryTypes.String) 1479 }) 1480 } 1481 }) 1482 } 1483 }) 1484 }