Skip to content

Commit 27cea43

Browse files
zeroshadezeroshade
authored
apacheGH-20415: [Go] Kernel Input Type for RLE (apache#14146)
Authored-by: Matt Topol <zotthewizard@gmail.com> Signed-off-by: Matt Topol <zotthewizard@gmail.com>
1 parent d422137 commit 27cea43

3 files changed

Lines changed: 69 additions & 1 deletion

File tree

‎go/arrow/array/concat.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ func Concatenate(arrs []arrow.Array, mem memory.Allocator) (result arrow.Array,
4747
case error:
4848
err = fmt.Errorf("arrow/concat: %w", e)
4949
default:
50-
err = fmt.Errorf("arrow/concat: unknown error: %v", pErr)
50+
err = fmt.Errorf("arrow/concat: %v", pErr)
5151
}
5252
}
5353
}()

‎go/arrow/compute/internal/exec/kernel.go‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -242,6 +242,40 @@ func (primitiveMatcher) Equals(other TypeMatcher) bool {
242242
// returns true for.
243243
func Primitive() TypeMatcher { return primitiveMatcher{} }
244244

245+
type reeMatcher struct {
246+
runEndsMatcher TypeMatcher
247+
encodedMatcher TypeMatcher
248+
}
249+
250+
func (r reeMatcher) Matches(typ arrow.DataType) bool {
251+
if typ.ID() != arrow.RUN_END_ENCODED {
252+
return false
253+
}
254+
255+
dt := typ.(*arrow.RunEndEncodedType)
256+
return r.runEndsMatcher.Matches(dt.RunEnds()) && r.encodedMatcher.Matches(dt.Encoded())
257+
}
258+
259+
func (r reeMatcher) Equals(other TypeMatcher) bool {
260+
o, ok := other.(reeMatcher)
261+
if !ok {
262+
return false
263+
}
264+
return r.runEndsMatcher.Equals(o.runEndsMatcher) && r.encodedMatcher.Equals(o.encodedMatcher)
265+
}
266+
267+
func (r reeMatcher) String() string {
268+
return "run_end_encoded(run_ends=" + r.runEndsMatcher.String() + ", values=" + r.encodedMatcher.String() + ")"
269+
}
270+
271+
// RunEndEncoded returns a matcher which matches a RunEndEncoded
272+
// type whose encoded type is matched by the passed in matcher.
273+
func RunEndEncoded(runEndsMatcher, encodedMatcher TypeMatcher) TypeMatcher {
274+
return reeMatcher{
275+
runEndsMatcher: runEndsMatcher,
276+
encodedMatcher: encodedMatcher}
277+
}
278+
245279
// InputKind is an enum representing the type of Input matching
246280
// that will be done. Either accepting any type, an exact specific type
247281
// or using a TypeMatcher.

‎go/arrow/compute/internal/exec/kernel_test.go‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -161,6 +161,40 @@ func TestPrimitiveMatcher(t *testing.T) {
161161
assert.False(t, match.Matches(arrow.Null))
162162
}
163163

164+
func TestREEMatcher(t *testing.T) {
165+
tests := []struct {
166+
runEnds exec.TypeMatcher
167+
enc exec.TypeMatcher
168+
matchRunEnds arrow.DataType
169+
nomatchRunEnds arrow.DataType
170+
matchEnc arrow.DataType
171+
nomatchEnc arrow.DataType
172+
}{
173+
{exec.Integer(), exec.Integer(), arrow.PrimitiveTypes.Int16, arrow.FixedWidthTypes.Float16, arrow.PrimitiveTypes.Int8, arrow.BinaryTypes.String},
174+
{exec.SameTypeID(arrow.INT32), exec.BinaryLike(), arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int64, arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int32},
175+
{exec.SameTypeID(arrow.INT64), exec.SameTypeID(arrow.STRUCT), arrow.PrimitiveTypes.Int64, arrow.PrimitiveTypes.Int32, arrow.StructOf(arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int16}), arrow.PrimitiveTypes.Int8},
176+
}
177+
178+
for _, tt := range tests {
179+
t.Run(tt.enc.String(), func(t *testing.T) {
180+
matcher := exec.RunEndEncoded(tt.runEnds, tt.enc)
181+
assert.False(t, matcher.Matches(tt.matchEnc))
182+
assert.True(t, matcher.Matches(arrow.RunEndEncodedOf(tt.matchRunEnds, tt.matchEnc)))
183+
assert.False(t, matcher.Matches(arrow.RunEndEncodedOf(tt.matchRunEnds, tt.nomatchEnc)))
184+
assert.False(t, matcher.Matches(arrow.RunEndEncodedOf(tt.nomatchRunEnds, tt.matchEnc)))
185+
assert.False(t, matcher.Matches(arrow.RunEndEncodedOf(tt.nomatchRunEnds, tt.nomatchEnc)))
186+
187+
assert.Equal(t, "run_end_encoded(run_ends="+tt.runEnds.String()+", values="+tt.enc.String()+")", matcher.String())
188+
189+
assert.True(t, matcher.Equals(exec.RunEndEncoded(tt.runEnds, tt.enc)))
190+
assert.False(t, matcher.Equals(exec.Primitive()))
191+
assert.False(t, matcher.Equals(exec.RunEndEncoded(exec.SameTypeID(tt.nomatchRunEnds.ID()), exec.SameTypeID(tt.nomatchEnc.ID()))))
192+
assert.False(t, matcher.Equals(exec.RunEndEncoded(exec.SameTypeID(tt.matchRunEnds.ID()), exec.SameTypeID(tt.nomatchEnc.ID()))))
193+
assert.False(t, matcher.Equals(exec.RunEndEncoded(exec.SameTypeID(tt.nomatchRunEnds.ID()), exec.SameTypeID(tt.matchEnc.ID()))))
194+
})
195+
}
196+
}
197+
164198
func TestInputTypeAnyType(t *testing.T) {
165199
var ty exec.InputType
166200
assert.Equal(t, exec.InputAny, ty.Kind)

0 commit comments

Comments
 (0)