Support arrow.TIME64 arrow.TIMESTAMP .. in exprs.ToSubstraitType()

[Sunyue]

Describe the enhancement requested

I am trying to do filter on arrow table/record with compute and exprs package, and found filtering for arrow.TIME64, arrow.TIMESTAMP, arrow.DECIMAL128, arrow.DECIMAL256 is not working as expected.

Could you please help to support arrow.TIME64 arrow.TIMESTAMP ... in exprs.ToSubstraitType()
https://pkg.go.dev/github.com/apache/arrow-go/v18/arrow/compute/exprs#ToSubstraitType

Looking at the code arrow.DECIMAL128, arrow.DECIMAL256 seems are supported, but not work well, maybe a bug?

Expected: SetInputSchema(), exprs.ExecuteScalarExpression() can work well with arrow.TIME64 arrow.TIMESTAMP, arrow.DECIMAL128, arrow.DECIMAL256

Note: I find it relatively simple to support arrow.TIME64 arrow.TIMESTAMP by adding cases to exprs.ToSubstraitType(), but it seems only arrow.Microsecond is supported as it is hard coded in FromSubstraitType, any possibility to derive the unit from the input schema?

Component(s)

[zeroshade]

I'll take a look at how difficult it'll be to update the conversions. If you have any code that you were using for testing, can you share it?

[Sunyue]

I created a test case like below, thanks again for your help!
(Other types like int/uint/float/string/binary/date32 seems working well with filter conditions like less, equal, ..., it will be great if time, timestamp, decimal can do the same)

import (
	"testing"

	"github.com/apache/arrow-go/v18/arrow"
	"github.com/apache/arrow-go/v18/arrow/array"
	"github.com/apache/arrow-go/v18/arrow/compute"
	"github.com/apache/arrow-go/v18/arrow/compute/exprs"
	"github.com/apache/arrow-go/v18/arrow/decimal"
	"github.com/apache/arrow-go/v18/arrow/memory"
	"github.com/stretchr/testify/require"
	"github.com/substrait-io/substrait-go/v3/expr"
	"github.com/substrait-io/substrait-go/v3/types"
)

func Test_Types(t *testing.T) {
	t.Parallel()

	tt := []struct {
		name   string
		schema func() *arrow.Schema
		record func(rq *require.Assertions, schema *arrow.Schema) arrow.Record
		val    func(rq *require.Assertions) expr.Literal
	}{
		{
			name: "expect arrow.TIME64 (ns) ok",
			schema: func() *arrow.Schema {
				field := arrow.Field{
					Name:     "col",
					Type:     &arrow.Time64Type{Unit: arrow.Nanosecond},
					Nullable: true,
				}

				return arrow.NewSchema([]arrow.Field{field}, nil)
			},
			record: func(rq *require.Assertions, schema *arrow.Schema) arrow.Record {
				b := array.NewTime64Builder(memory.DefaultAllocator, &arrow.Time64Type{Unit: arrow.Nanosecond})
				defer b.Release()

				t1, err := arrow.Time64FromString("10:00:00.000000", arrow.Nanosecond)
				rq.NoError(err, "Failed to create Time64 value")

				b.AppendValues([]arrow.Time64{t1}, []bool{true})

				return array.NewRecord(schema, []arrow.Array{b.NewArray()}, 1)
			},
			val: func(rq *require.Assertions) expr.Literal {
				v, err := arrow.Time64FromString("11:00:00.000000", arrow.Nanosecond)
				rq.NoError(err, "Failed to create Time64 value")

				return expr.NewPrimitiveLiteral(types.Time(v), true)
			},
		},
		{
			name: "expect arrow.TIMESTAMP (ns) ok",
			schema: func() *arrow.Schema {
				field := arrow.Field{
					Name:     "col",
					Type:     &arrow.TimestampType{Unit: arrow.Nanosecond},
					Nullable: true,
				}

				return arrow.NewSchema([]arrow.Field{field}, nil)
			},
			record: func(rq *require.Assertions, schema *arrow.Schema) arrow.Record {
				b := array.NewTimestampBuilder(memory.DefaultAllocator, &arrow.TimestampType{Unit: arrow.Nanosecond})
				defer b.Release()

				t1, err := arrow.TimestampFromString("2021-01-01T10:00:00.000000Z", arrow.Nanosecond)
				rq.NoError(err, "Failed to create Timestamp value")

				b.AppendValues([]arrow.Timestamp{t1}, []bool{true})

				return array.NewRecord(schema, []arrow.Array{b.NewArray()}, 1)
			},
			val: func(rq *require.Assertions) expr.Literal {
				v, err := arrow.TimestampFromString("2021-01-01T11:00:00.000000Z", arrow.Nanosecond)
				rq.NoError(err, "Failed to create Timestamp value")

				return expr.NewPrimitiveLiteral(types.Timestamp(v), true)
			},
		},
		{
			name: "expect arrow.DECIMAL128 ok",
			schema: func() *arrow.Schema {
				field := arrow.Field{
					Name:     "col",
					Type:     &arrow.Decimal128Type{Precision: 38, Scale: 10},
					Nullable: true,
				}

				return arrow.NewSchema([]arrow.Field{field}, nil)
			},
			record: func(rq *require.Assertions, schema *arrow.Schema) arrow.Record {
				b := array.NewDecimal128Builder(memory.DefaultAllocator, &arrow.Decimal128Type{Precision: 38, Scale: 10})
				defer b.Release()

				d, err := decimal.Decimal128FromFloat(123.456789, 38, 10)
				rq.NoError(err, "Failed to create Decimal128 value")

				b.Append(d)

				return array.NewRecord(schema, []arrow.Array{b.NewArray()}, 1)
			},
			val: func(rq *require.Assertions) expr.Literal {
				v, p, s, err := expr.DecimalStringToBytes("456.7890123456")
				rq.NoError(err, "Failed to convert decimal string to bytes")

				lit, err := expr.NewLiteral(&types.Decimal{
					Value:     v[:16],
					Precision: p,
					Scale:     s,
				}, true)
				rq.NoError(err, "Failed to create Decimal128 literal")

				return lit
			},
		},
	}

	for _, tc := range tt {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			ctx := t.Context()
			rq := require.New(t)
			schema := tc.schema()
			record := tc.record(rq, schema)

			extSet := exprs.GetExtensionIDSet(ctx)
			builder := exprs.NewExprBuilder(extSet)

			err := builder.SetInputSchema(schema)
			rq.NoError(err, "Failed to set input schema")

			b, err := builder.CallScalar("less", nil,
				builder.FieldRef("col"),
				builder.Literal(tc.val(rq)),
			)

			rq.NoError(err, "Failed to call scalar")

			e, err := b.BuildExpr()
			rq.NoError(err, "Failed to build expression")

			ctx = exprs.WithExtensionIDSet(ctx, extSet)

			dr := compute.NewDatum(record)
			defer dr.Release()

			_, err = exprs.ExecuteScalarExpression(ctx, schema, e, dr)
			rq.NoError(err, "Failed to execute scalar expression")
		})
	}
}

[Sunyue]

Hi @zeroshade, I had a try with this enhancement using v18.4.0-rc0, when executing exprs.ExecuteScalarExpression it still gives me an unimplemented error. The error is raised from arrow/compute/exprs/exec.go literalToDatum(), maybe those new PrecisionTimeType or PrecisionTimestampType should also be supported in this method? Thanks.

[Sunyue]

Please find my test case below:

import (
	"testing"

	"github.com/apache/arrow-go/v18/arrow"
	"github.com/apache/arrow-go/v18/arrow/array"
	"github.com/apache/arrow-go/v18/arrow/compute"
	"github.com/apache/arrow-go/v18/arrow/compute/exprs"
	"github.com/apache/arrow-go/v18/arrow/memory"
	"github.com/stretchr/testify/require"
	"github.com/substrait-io/substrait-go/v4/expr"
	"github.com/substrait-io/substrait-go/v4/types"
)

func Test_Types(t *testing.T) {
	t.Parallel()

	tt := []struct {
		name   string
		schema func() *arrow.Schema
		record func(rq *require.Assertions, schema *arrow.Schema) arrow.Record
		val    func(rq *require.Assertions) expr.Literal
	}{
		{
			name: "expect arrow.TIME64 (ns) ok",
			schema: func() *arrow.Schema {
				field := arrow.Field{
					Name:     "col",
					Type:     &arrow.Time64Type{Unit: arrow.Nanosecond},
					Nullable: true,
				}

				return arrow.NewSchema([]arrow.Field{field}, nil)
			},
			record: func(rq *require.Assertions, schema *arrow.Schema) arrow.Record {
				b := array.NewTime64Builder(memory.DefaultAllocator, &arrow.Time64Type{Unit: arrow.Nanosecond})
				defer b.Release()

				t1, err := arrow.Time64FromString("10:00:00.000000", arrow.Nanosecond)
				rq.NoError(err, "Failed to create Time64 value")

				b.AppendValues([]arrow.Time64{t1}, []bool{true})

				return array.NewRecord(schema, []arrow.Array{b.NewArray()}, 1)
			},
			val: func(rq *require.Assertions) expr.Literal {
				v, err := arrow.Time64FromString("11:00:00.000000", arrow.Nanosecond)
				rq.NoError(err, "Failed to create Time64 value")

				pt := &types.PrecisionTime{
					Precision: int32(arrow.Nanosecond) * 3,
					Value:     int64(v),
				}

				lit, err := expr.NewLiteral(pt, true)
				rq.NoError(err, "Failed to create literal")

				return lit
			},
		},
	}

	for _, tc := range tt {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			ctx := t.Context()
			rq := require.New(t)
			schema := tc.schema()
			record := tc.record(rq, schema)

			extSet := exprs.GetExtensionIDSet(ctx)
			builder := exprs.NewExprBuilder(extSet)

			err := builder.SetInputSchema(schema)
			rq.NoError(err, "Failed to set input schema")

			b, err := builder.CallScalar("less", nil,
				builder.FieldRef("col"),
				builder.Literal(tc.val(rq)),
			)

			rq.NoError(err, "Failed to call scalar")

			e, err := b.BuildExpr()
			rq.NoError(err, "Failed to build expression")

			ctx = exprs.WithExtensionIDSet(ctx, extSet)

			dr := compute.NewDatum(record)
			defer dr.Release()

			_, err = exprs.ExecuteScalarExpression(ctx, schema, e, dr)
			rq.NoError(err, "Failed to execute scalar expression")
		})
	}
}

[Sunyue]

Another thing is filtering on arrow.DECIMAL128 now works, but still not the case for arrow.DECIMAL256. It looks like when converting to types.DecimalType, both DECIMAL128 and DECIMAL256 are possible, but there are several place when converting from types.DecimalType back to arrow types, it was hard coded to DECIMAL128. Any chance make it work for DECIMAL256?