diff --git a/df/arrow/df.go b/df/arrow/df.go new file mode 100644 index 0000000..42d97a3 --- /dev/null +++ b/df/arrow/df.go @@ -0,0 +1,1167 @@ +//go:build arrow + +package arrow + +import ( + "context" + "fmt" + "reflect" + "strings" + "time" + + "github.com/apache/arrow/go/v14/arrow" + "github.com/apache/arrow/go/v14/arrow/array" + "github.com/apache/arrow/go/v14/arrow/builder" + "github.com/apache/arrow/go/v14/arrow/compute" + "github.com/apache/arrow/go/v14/arrow/memory" + "github.com/apache/arrow/go/v14/arrow/scalar" + "git.querycap.com/practice/df" // MODIFIED import path +) + +type arrowDataFrame struct { + name string + schema *arrowDataFrameSchema + record arrow.Record + mem memory.Allocator +} + +func NewArrowDataFrame(name string, record arrow.Record, dfSchema *arrowDataFrameSchema) df.DataFrame { + return NewArrowDataFrameWithAllocator(name, record, dfSchema, memory.DefaultAllocator) +} +func NewArrowDataFrameWithAllocator(name string, record arrow.Record, dfSchema *arrowDataFrameSchema, mem memory.Allocator) df.DataFrame { + if dfSchema == nil { panic("NewArrowDataFrameWithAllocator: df.DataFrameSchema cannot be nil") } + if mem == nil { panic("NewArrowDataFrameWithAllocator: memory.Allocator cannot be nil") } + if record == nil { + if !(dfSchema.schema == nil || dfSchema.schema.NumFields() == 0) { + panic("NewArrowDataFrameWithAllocator: record is nil but schema defines fields") + } + } else { + if dfSchema.schema == nil { panic("NewArrowDataFrameWithAllocator: dfSchema.schema is nil for a non-nil record") } + if !dfSchema.schema.Equal(record.Schema()) { + panic(fmt.Sprintf("NewArrowDataFrameWithAllocator: schema mismatch. Provided: %s, Record: %s", dfSchema.schema, record.Schema())) + } + record.Retain() + } + return &arrowDataFrame{name: name, schema: dfSchema, record: record, mem: mem} +} +func NewArrowDataFrameFromArrays(name string, cols []arrow.Array, schema *arrow.Schema) (df.DataFrame, error) { + return NewArrowDataFrameFromArraysWithAllocator(name, cols, schema, memory.DefaultAllocator) +} +func NewArrowDataFrameFromArraysWithAllocator(name string, cols []arrow.Array, schema *arrow.Schema, mem memory.Allocator) (df.DataFrame, error) { + if schema == nil {return nil, fmt.Errorf("arrow.Schema cannot be nil")} + if mem == nil {return nil, fmt.Errorf("memory.Allocator cannot be nil")} + if len(cols) != schema.NumFields() {return nil, fmt.Errorf("num cols (%d) != num fields (%d)", len(cols), schema.NumFields())} + var numRows int64 = -1 + if len(cols) > 0 { + for i := range cols { cols[i].Retain() } + numRows = int64(cols[0].Len()) + for i, col := range cols { + if int64(col.Len()) != numRows { + for j := 0; j <= i; j++ { cols[j].Release() }; return nil, fmt.Errorf("col %d len %d != %d", i, col.Len(), numRows) + } + if !arrow.TypeEqual(col.DataType(), schema.Field(i).Type) { + for j := 0; j <= i; j++ { cols[j].Release() }; return nil, fmt.Errorf("col %d type %s != schema %s", i, col.DataType(), schema.Field(i).Type) + } + } + } else {numRows = 0} + record := array.NewRecord(schema, cols, numRows); + for _, col := range cols { col.Release() } // NewRecord created its own references or copied data. + dfSchema := NewArrowDataFrameSchema(schema).(*arrowDataFrameSchema) + defer record.Release() // Release the record created here after NewArrowDataFrameWithAllocator is done. + return NewArrowDataFrameWithAllocator(name, record, dfSchema, mem), nil +} +func NewArrowDataFrameFromSeries(name string, series []df.Series, mem memory.Allocator) (df.DataFrame, error) { + if mem == nil { mem = memory.DefaultAllocator } + if len(series) == 0 { + emptyArrowSchema := arrow.NewSchema([]arrow.Field{}, nil) + emptyDfSchema := NewArrowDataFrameSchema(emptyArrowSchema).(*arrowDataFrameSchema) + emptyRecord := array.NewRecord(emptyArrowSchema, nil, 0) + return NewArrowDataFrameWithAllocator(name, emptyRecord, emptyDfSchema, mem), nil + } + arrowArrays := make([]arrow.Array, len(series)) + arrowFields := make([]arrow.Field, len(series)) + var numRows int = -1 + for i, s := range series { + as, ok := s.(*arrowSeries); + if !ok { for j := 0; j < i; j++ { arrowArrays[j].Release() }; return nil, fmt.Errorf("NewArrowDataFrameFromSeries: all series must be *arrowSeries, found %T at index %d", s, i) } + if numRows == -1 { numRows = as.Len() } else if as.Len() != numRows { for j := 0; j < i; j++ { arrowArrays[j].Release() }; return nil, fmt.Errorf("NewArrowDataFrameFromSeries: series length mismatch") } + as.arr.Retain(); arrowArrays[i] = as.arr + sSchema := as.Schema(); arrowDataType, err := dfFormatToArrowType(sSchema.Format) + if err != nil { for j := 0; j <=i; j++ { arrowArrays[j].Release()}; return nil, fmt.Errorf("NewArrowDataFrameFromSeries: %w", err)} + arrowFields[i] = arrow.Field{Name: sSchema.Name, Type: arrowDataType, Nullable: sSchema.Nullable, Metadata: arrow.MetadataFrom(sSchema.Metadata)} + } + arrowSchema := arrow.NewSchema(arrowFields, nil) + record := array.NewRecord(arrowSchema, arrowArrays, int64(numRows)) + for _, arr := range arrowArrays { arr.Release() } // Record has them now + dfSchema := NewArrowDataFrameSchema(record.Schema()).(*arrowDataFrameSchema) + defer record.Release() + return NewArrowDataFrameWithAllocator(name, record, dfSchema, mem), nil +} + +func (adf *arrowDataFrame) Schema() df.DataFrameSchema { return adf.schema } +func (adf *arrowDataFrame) Name() string { return adf.name } +func (adf *arrowDataFrame) Len() int { if adf.record == nil { return 0 }; return int(adf.record.NumRows()) } +func (adf *arrowDataFrame) Release() { if adf.record != nil { adf.record.Release(); adf.record = nil } } +func (adf *arrowDataFrame) GetSeries(index int) df.Series { + if adf.record == nil || index < 0 || index >= int(adf.record.NumCols()) { panic(fmt.Sprintf("series index %d out of bounds", index)) } + return NewArrowSeriesWithAllocator(adf.record.Column(index), adf.schema.Get(index), adf.mem) +} +func (adf *arrowDataFrame) GetSeriesByName(sName string) df.Series { + idx := adf.schema.GetIndexByName(sName); if idx == -1 { panic(fmt.Sprintf("series '%s' not found", sName)) }; return adf.GetSeries(idx) +} +func (adf *arrowDataFrame) GetRow(i int64) df.Row { + if adf.record == nil || i < 0 || i >= adf.record.NumRows() { panic(fmt.Sprintf("row index %d out of bounds", i)) } + r, err := NewArrowRowFromRecord(adf.schema, adf.record, int(i)); if err != nil { panic(err) }; return r +} +func (adf *arrowDataFrame) GetValue(rowIndx, colIndx int) df.Value { + if adf.record == nil || rowIndx < 0 || int64(rowIndx) >= adf.record.NumRows() || colIndx < 0 || colIndx >= int(adf.record.NumCols()) { + panic(fmt.Sprintf("GetValue index (row: %d, col: %d) out of bounds", rowIndx, colIndx)) + } + return NewArrowValue(scalar.MakeScalar(adf.record.Column(colIndx), rowIndx), adf.schema.Get(colIndx).Format) +} +func (adf *arrowDataFrame) Limit(offset int, size int) df.DataFrame { + if adf.record == nil { emptyRec := array.NewRecord(adf.schema.schema, nil, 0); defer emptyRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, emptyRec, adf.schema, adf.mem) } + currentNumRows := adf.record.NumRows(); if offset < 0 { offset = 0 } + if offset >= int(currentNumRows) { emptyRec := array.NewRecord(adf.schema.schema, nil, 0); defer emptyRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, emptyRec, adf.schema, adf.mem) } + if offset+size > int(currentNumRows) { size = int(currentNumRows) - offset } + if size <= 0 { emptyRec := array.NewRecord(adf.schema.schema, nil, 0); defer emptyRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, emptyRec, adf.schema, adf.mem) } + slicedRecord := adf.record.NewSlice(int64(offset), int64(offset+size)); defer slicedRecord.Release() + return NewArrowDataFrameWithAllocator(adf.name, slicedRecord, adf.schema, adf.mem) +} +func (adf *arrowDataFrame) SelectBySeriesIndex(indices ...int) df.DataFrame { + numRowsToKeep := int64(0); if adf.record != nil { numRowsToKeep = adf.record.NumRows()} + if len(indices) == 0 { + emptyArrowSchema := arrow.NewSchema([]arrow.Field{}, nil); emptyDfSchema := NewArrowDataFrameSchema(emptyArrowSchema).(*arrowDataFrameSchema) + emptyRecord := array.NewRecord(emptyArrowSchema, nil, numRowsToKeep); defer emptyRecord.Release() + return NewArrowDataFrameWithAllocator(adf.name, emptyRecord, emptyDfSchema, adf.mem) + } + if adf.record == nil { panic("cannot select columns from a nil or released dataframe") } + newFields := make([]arrow.Field, len(indices)); newCols := make([]arrow.Array, len(indices)) + for i, idx := range indices { + if idx < 0 || idx >= int(adf.record.NumCols()) { for j := 0; j < i; j++ { newCols[j].Release() }; panic(fmt.Sprintf("select index %d out of bounds", idx)) } + newFields[i] = adf.schema.schema.Field(idx); newCols[i] = adf.record.Column(idx); newCols[i].Retain() + } + newArrowSchema := arrow.NewSchema(newFields, adf.schema.schema.Metadata()); newDfSchema := NewArrowDataFrameSchema(newArrowSchema).(*arrowDataFrameSchema) + selectedRecord := array.NewRecord(newArrowSchema, newCols, numRowsToKeep) + for _, col := range newCols { col.Release() }; defer selectedRecord.Release() + return NewArrowDataFrameWithAllocator(adf.name, selectedRecord, newDfSchema, adf.mem) +} +func (adf *arrowDataFrame) SelectBySeriesName(colNames ...string) df.DataFrame { + if adf.record == nil && len(colNames) > 0 { panic("cannot select by name from a nil or released dataframe") } + if len(colNames) == 0 { + numRowsToKeep := int64(0); if adf.record != nil { numRowsToKeep = adf.record.NumRows()} + emptyArrowSchema := arrow.NewSchema([]arrow.Field{}, nil); emptyDfSchema := NewArrowDataFrameSchema(emptyArrowSchema).(*arrowDataFrameSchema) + emptyRecord := array.NewRecord(emptyArrowSchema, nil, numRowsToKeep); defer emptyRecord.Release() + return NewArrowDataFrameWithAllocator(adf.name, emptyRecord, emptyDfSchema, adf.mem) + } + indices := make([]int, len(colNames)) + for i, name := range colNames { idx := adf.schema.GetIndexByName(name); if idx == -1 { panic(fmt.Sprintf("column '%s' not found", name)) }; indices[i] = idx } + return adf.SelectBySeriesIndex(indices...) +} +func (adf *arrowDataFrame) WhereRow(f func(df.Row) bool) df.DataFrame { + if adf.record == nil { emptyRec := array.NewRecord(adf.schema.schema, nil, 0); defer emptyRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, emptyRec, adf.schema, adf.mem) } + numCols := int(adf.record.NumCols()); currentSchema := adf.schema.schema + colBuilders := make([]array.Builder, numCols); for i := 0; i < numCols; i++ { colBuilders[i] = builder.NewBuilder(adf.mem, currentSchema.Field(i).Type) } + defer func() { for _, b := range colBuilders { if b != nil { b.Release() } } }() + for r := int64(0); r < adf.record.NumRows(); r++ { + rowView, _ := NewArrowRowFromRecord(adf.schema, adf.record, int(r)) + if f(rowView) { + for c := 0; c < numCols; c++ { if err := array.CopyValue(colBuilders[c], adf.record.Column(c), int(r)); err != nil { panic(fmt.Sprintf("error copying value col %d, row %d: %v", c,r,err))}} + } + } + newCols := make([]arrow.Array, numCols); var newRecordLen int64 + if len(colBuilders) > 0 && colBuilders[0] != nil { newRecordLen = int64(colBuilders[0].Len()) } else { newRecordLen = 0 } + for i, b := range colBuilders { newCols[i] = b.NewArray() } + filteredRecord := array.NewRecord(currentSchema, newCols, newRecordLen) + for _, col := range newCols { col.Release() }; defer filteredRecord.Release() + return NewArrowDataFrameWithAllocator(adf.name, filteredRecord, adf.schema, adf.mem) +} +func (adf *arrowDataFrame) Sort(orders ...df.SortByIndex) df.DataFrame { + if adf.record == nil || adf.record.NumRows() == 0 || len(orders) == 0 { + var recToHandle arrow.Record + if adf.record != nil { recToHandle = adf.record.NewSlice(0, adf.record.NumRows()) } else { recToHandle = array.NewRecord(adf.schema.schema, nil, 0) } + defer recToHandle.Release(); return NewArrowDataFrameWithAllocator(adf.name, recToHandle, adf.schema, adf.mem) + } + ctx := compute.WithAllocator(context.Background(), adf.mem); sortKeys := make([]compute.SortKey, len(orders)) + for i, order := range orders { + if order.Series < 0 || order.Series >= int(adf.record.NumCols()) { panic(fmt.Sprintf("sort key index %d out of bounds", order.Series)) } + arrowSortOrder := arrow.Ascending; if order.Order == df.SortOrderDESC { arrowSortOrder = arrow.Descending } + sortKeys[i] = compute.SortKey{ Name: adf.schema.schema.Field(order.Series).Name, Order: arrowSortOrder } + } + indicesDatum, err := compute.SortIndices(ctx, arrow.NewRecordDatum(adf.record), compute.SortOptions{SortKeys: sortKeys, NullPlacement: arrow.NullsFirst}) + if err != nil { panic(fmt.Sprintf("failed to get sort indices: %v", err)) }; defer indicesDatum.Release() + indicesArr, ok := indicesDatum.(*arrow.ArrayDatum).Value.(arrow.Array); if !ok { panic("SortIndices bad return type") } + sortedRecordDatum, err := compute.Take(ctx, compute.TakeOptions{}, arrow.NewRecordDatum(adf.record), arrow.NewArrayDatum(indicesArr)) + if err != nil { panic(fmt.Sprintf("failed to take sorted rows: %v", err)) }; defer sortedRecordDatum.Release() + sortedRecord, ok := sortedRecordDatum.(*arrow.RecordDatum).Value().(arrow.Record); if !ok { panic("Take bad return type") } + return NewArrowDataFrameWithAllocator(adf.name, sortedRecord, adf.schema, adf.mem) +} +func (adf *arrowDataFrame) SortByName(orders ...df.SortByName) df.DataFrame { + if adf.record == nil && len(orders) > 0 { panic("cannot sort by name on nil dataframe") } + if len(orders) == 0 { + var recToHandle arrow.Record + if adf.record != nil { recToHandle = adf.record.NewSlice(0, adf.record.NumRows()) } else { recToHandle = array.NewRecord(adf.schema.schema, nil, 0) } + defer recToHandle.Release(); return NewArrowDataFrameWithAllocator(adf.name, recToHandle, adf.schema, adf.mem) + } + sortByIdx := make([]df.SortByIndex, len(orders)) + for i, order := range orders { idx := adf.schema.GetIndexByName(order.Series); if idx == -1 { panic(fmt.Sprintf("col '%s' not found", order.Series)) }; sortByIdx[i] = df.SortByIndex{Series: idx, Order: order.Order} } + return adf.Sort(sortByIdx...) +} +func (adf *arrowDataFrame) AddSeries(colName string, series df.Series) df.DataFrame { + if adf.record == nil { panic("nil dataframe record") }; if adf.schema.HasName(colName) { panic(fmt.Sprintf("col '%s' exists", colName)) } + arrowSeries, ok := series.(*arrowSeries); if !ok { panic(fmt.Sprintf("expected *arrowSeries, got %T", series)) } + if arrowSeries.arr == nil { panic("new series array is nil") }; if arrowSeries.Len() != adf.Len() { panic(fmt.Sprintf("len mismatch: df %d, series %d", adf.Len(), arrowSeries.Len())) } + existingFields := adf.schema.schema.Fields(); newSchemaFields := make([]arrow.Field, len(existingFields)+1) + copy(newSchemaFields, existingFields) + newSchemaFields[len(existingFields)] = arrow.Field{Name: colName, Type: arrowSeries.arr.DataType(), Nullable: arrowSeries.arr.NullN() > 0} + newArrowSchema := arrow.NewSchema(newSchemaFields, adf.schema.schema.Metadata()); newDfSchema := NewArrowDataFrameSchema(newArrowSchema).(*arrowDataFrameSchema) + existingCols := adf.record.Columns(); newRecordCols := make([]arrow.Array, len(existingCols)+1) + for i, col := range existingCols { col.Retain(); newRecordCols[i] = col } + arrowSeries.arr.Retain(); newRecordCols[len(existingCols)] = arrowSeries.arr + newRecord := array.NewRecord(newArrowSchema, newRecordCols, adf.record.NumRows()) + for _, col := range newRecordCols { col.Release() }; defer newRecord.Release() + return NewArrowDataFrameWithAllocator(adf.name, newRecord, newDfSchema, adf.mem) +} +func (adf *arrowDataFrame) RemoveSeries(index int) df.DataFrame { + if adf.record == nil { panic("nil dataframe record") }; if index < 0 || index >= int(adf.record.NumCols()) { panic(fmt.Sprintf("index %d out of bounds", index)) } + numOldCols := int(adf.record.NumCols()); newSchemaFields := make([]arrow.Field,0,numOldCols-1); newRecordCols := make([]arrow.Array,0,numOldCols-1) + for i := 0; i < numOldCols; i++ { + if i == index { continue }; newSchemaFields = append(newSchemaFields, adf.schema.schema.Field(i)) + col := adf.record.Column(i); col.Retain(); newRecordCols = append(newRecordCols, col) + } + newArrowSchema := arrow.NewSchema(newSchemaFields, adf.schema.schema.Metadata()); newDfSchema := NewArrowDataFrameSchema(newArrowSchema).(*arrowDataFrameSchema) + newRecord := array.NewRecord(newArrowSchema, newRecordCols, adf.record.NumRows()) + for _, col := range newRecordCols { col.Release() }; defer newRecord.Release() + return NewArrowDataFrameWithAllocator(adf.name, newRecord, newDfSchema, adf.mem) +} +func (adf *arrowDataFrame) RemoveSeriesByName(s string) df.DataFrame { + idx := adf.schema.GetIndexByName(s); if idx == -1 { panic(fmt.Sprintf("col '%s' not found", s))}; return adf.RemoveSeries(idx) +} +func (adf *arrowDataFrame) RenameSeries(index int, newName string, inplace bool) df.DataFrame { + if adf.record == nil { panic("nil dataframe record") }; if index < 0 || index >= int(adf.record.NumCols()) { panic(fmt.Sprintf("index %d out of bounds", index)) } + if currentName := adf.schema.schema.Field(index).Name; currentName == newName { + if inplace { return adf }; newRecView := adf.record.NewSlice(0, adf.record.NumRows()); defer newRecView.Release(); return NewArrowDataFrameWithAllocator(adf.name, newRecView, adf.schema, adf.mem) + } + if adf.schema.HasName(newName) { panic(fmt.Sprintf("col '%s' exists", newName)) } + newSchemaFields := make([]arrow.Field, adf.record.NumCols()); copy(newSchemaFields, adf.schema.schema.Fields()) + newSchemaFields[index].Name = newName + newArrowSchema := arrow.NewSchema(newSchemaFields, adf.schema.schema.Metadata()); newDfSchema := NewArrowDataFrameSchema(newArrowSchema).(*arrowDataFrameSchema) + recordCols := adf.record.Columns(); for _, col := range recordCols { col.Retain() } + newRecord := array.NewRecord(newArrowSchema, recordCols, adf.record.NumRows()); for _, col := range recordCols { col.Release() } + if inplace { adf.record.Release(); adf.record = newRecord; adf.schema = newDfSchema; return adf } + defer newRecord.Release(); return NewArrowDataFrameWithAllocator(adf.name, newRecord, newDfSchema, adf.mem) +} +func (adf *arrowDataFrame) RenameSeriesByName(colName string, newName string, inplace bool) df.DataFrame { + idx := adf.schema.GetIndexByName(colName); if idx == -1 { panic(fmt.Sprintf("col '%s' not found", colName)) }; return adf.RenameSeries(idx, newName, inplace) +} +func (adf *arrowDataFrame) GetSeriesExprByName(sName string) df.Expr { + idx := adf.schema.GetIndexByName(sName); if idx == -1 { panic(fmt.Sprintf("series '%s' not found", sName)) } + seriesSchema := adf.schema.Get(idx) + switch seriesSchema.Format.Name() { + case df.BoolFormat.Name(): return df.NewBoolColExpr(sName) + case df.IntegerFormat.Name(): return df.NewIntColExpr(sName) + case df.DoubleFormat.Name(): return df.NewDoubleColExpr(sName) + case df.StringFormat.Name(): return df.NewStringColExpr(sName) + case df.DateTimeFormat.Name(): return df.NewDatetimeColExpr(sName) + default: panic(fmt.Sprintf("GetSeriesExprByName unsupported format: %s", seriesSchema.Format.Name())) + } +} +func (adf *arrowDataFrame) MapRow(outputSchemaGiven df.DataFrameSchema, f func(df.Row) df.Row) df.DataFrame { + if adf.record == nil { panic("MapRow on nil record") }; outputArrowDFSchema, ok := outputSchemaGiven.(*arrowDataFrameSchema); if !ok { panic(fmt.Sprintf("outputSchema must be *arrowDataFrameSchema, got %T", outputSchemaGiven)) } + outputInternalArrowSchema := outputArrowDFSchema.schema; if outputInternalArrowSchema == nil { panic("outputSchema internal schema is nil") } + numOutputCols := outputInternalArrowSchema.NumFields(); colBuilders := make([]array.Builder, numOutputCols) + for i := 0; i < numOutputCols; i++ { colBuilders[i] = builder.NewBuilder(adf.mem, outputInternalArrowSchema.Field(i).Type) } + defer func() { for _, b := range colBuilders { if b != nil { b.Release() } } }() + for r := int64(0); r < adf.record.NumRows(); r++ { + inputRow, err := NewArrowRowFromRecord(adf.schema, adf.record, int(r)); if err != nil { panic(fmt.Sprintf("MapRow create input row %d: %v", r, err)) } + outputRow := f(inputRow); if outputRow == nil { panic(fmt.Sprintf("MapRow func returned nil df.Row for input %d", r)) } + if outputRow.Len() != numOutputCols { panic(fmt.Sprintf("MapRow func returned %d cols, expected %d", outputRow.Len(), numOutputCols)) } + for c := 0; c < numOutputCols; c++ { + val := outputRow.Get(c) + if val == nil || val.IsNil() { colBuilders[c].AppendNull(); continue } + arrowVal, castOk := val.(*arrowValue); if !castOk { panic(fmt.Sprintf("MapRow func value col %d type %T, expected *arrowValue", c, val)) } + // Using 3-argument appendScalarToBuilder from types.go + if err := appendScalarToBuilder(colBuilders[c], arrowVal.val, colBuilders[c].Type()); err != nil { panic(fmt.Sprintf("MapRow append col %d (name: %s): %v. Scalar: %s, Builder: %s",c, outputInternalArrowSchema.Field(c).Name, err, arrowVal.val.DataType().Name(), colBuilders[c].Type().Name()))} + } + } + newCols := make([]arrow.Array, numOutputCols); var newRecordLen int64 + if len(colBuilders) > 0 && colBuilders[0] != nil { newRecordLen = int64(colBuilders[0].Len()) } + for i, b := range colBuilders { newCols[i] = b.NewArray() } + mappedRecord := array.NewRecord(outputInternalArrowSchema, newCols, newRecordLen) + for _, col := range newCols { col.Release() }; defer mappedRecord.Release() + return NewArrowDataFrameWithAllocator(adf.name, mappedRecord, outputArrowDFSchema, adf.mem) +} +func (adf *arrowDataFrame) FlatMapRow(outputSchemaGiven df.DataFrameSchema, f func(df.Row) []df.Row) df.DataFrame { + if adf.record == nil { panic("FlatMapRow on nil record") }; outputArrowDFSchema, ok := outputSchemaGiven.(*arrowDataFrameSchema); if !ok { panic(fmt.Sprintf("outputSchema must be *arrowDataFrameSchema, got %T", outputSchemaGiven)) } + outputInternalArrowSchema := outputArrowDFSchema.schema; if outputInternalArrowSchema == nil { panic("outputSchema internal schema is nil") } + numOutputCols := outputInternalArrowSchema.NumFields(); colBuilders := make([]array.Builder, numOutputCols) + for i := 0; i < numOutputCols; i++ { colBuilders[i] = builder.NewBuilder(adf.mem, outputInternalArrowSchema.Field(i).Type) } + defer func() { for _, b := range colBuilders { if b != nil { b.Release() } } }() + for r := int64(0); r < adf.record.NumRows(); r++ { + inputRow, err := NewArrowRowFromRecord(adf.schema, adf.record, int(r)); if err != nil { panic(fmt.Sprintf("FlatMapRow create input row %d: %v", r, err)) } + outputRows := f(inputRow); if outputRows == nil { continue } + for i, outputRow := range outputRows { + if outputRow == nil { panic(fmt.Sprintf("FlatMapRow func returned slice with nil df.Row (index %d) for input %d", i, r)) } + if outputRow.Len() != numOutputCols { panic(fmt.Sprintf("FlatMapRow func returned df.Row (index %d) with %d cols, expected %d, for input %d", i, outputRow.Len(), numOutputCols, r)) } + for c := 0; c < numOutputCols; c++ { + val := outputRow.Get(c) + if val == nil || val.IsNil() { colBuilders[c].AppendNull(); continue } + arrowVal, castOk := val.(*arrowValue); if !castOk { panic(fmt.Sprintf("FlatMapRow func value (col %d, row %d) type %T, expected *arrowValue, for input %d", c, i, val, r)) } + // Using 3-argument appendScalarToBuilder from types.go + if err := appendScalarToBuilder(colBuilders[c], arrowVal.val, colBuilders[c].Type()); err != nil { panic(fmt.Sprintf("FlatMapRow append col %d (name: %s): %v. Scalar: %s, Builder: %s",c, outputInternalArrowSchema.Field(c).Name, err, arrowVal.val.DataType().Name(), colBuilders[c].Type().Name()))} + } + } + } + newCols := make([]arrow.Array, numOutputCols); var newRecordLen int64 + if len(colBuilders) > 0 && colBuilders[0] != nil { newRecordLen = int64(colBuilders[0].Len()) } + for i, b := range colBuilders { newCols[i] = b.NewArray() } + flatMappedRecord := array.NewRecord(outputInternalArrowSchema, newCols, newRecordLen) + for _, col := range newCols { col.Release() }; defer flatMappedRecord.Release() + return NewArrowDataFrameWithAllocator(adf.name, flatMappedRecord, outputArrowDFSchema, adf.mem) +} +func (adf *arrowDataFrame) Distinct(cols ...string) df.DataFrame { + if adf.record == nil || adf.record.NumRows() == 0 { schemaToUse := arrow.NewSchema([]arrow.Field{},nil); if adf.schema != nil && adf.schema.schema != nil { schemaToUse = adf.schema.schema}; newRec := array.NewRecord(schemaToUse, nil, 0); defer newRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, newRec, adf.schema, adf.mem)} + var keyIndices []int + if len(cols) == 0 { keyIndices = make([]int, adf.record.NumCols()); for i := 0; i < int(adf.record.NumCols()); i++ { keyIndices[i] = i } + } else { keyIndices = make([]int, len(cols)); for i, name := range cols { idx := adf.schema.GetIndexByName(name); if idx == -1 { panic(fmt.Sprintf("Distinct col '%s' not found", name)) }; keyIndices[i] = idx }} + if adf.record.NumCols() == 0 { if adf.record.NumRows() > 0 { newRec := adf.record.NewSlice(0,1); defer newRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, newRec, adf.schema, adf.mem) }; newRec := adf.record.NewSlice(0,0); defer newRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, newRec, adf.schema, adf.mem)} + sortOrders := make([]df.SortByIndex, len(keyIndices)); for i, keyIdx := range keyIndices { sortOrders[i] = df.SortByIndex{Series: keyIdx, Order: df.SortOrderASC} } + sortedDf := adf.Sort(sortOrders...); sortedArrowDf, ok := sortedDf.(*arrowDataFrame); if !ok { panic("Distinct: Sort bad return") }; defer sortedArrowDf.Release() + sortedRecord := sortedArrowDf.record; if sortedRecord.NumRows() == 0 { newRec := sortedRecord.NewSlice(0, 0); defer newRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, newRec, adf.schema, adf.mem)} + uniqueRowIndices := make([]int64, 0, sortedRecord.NumRows()); uniqueRowIndices = append(uniqueRowIndices, 0) + for i := int64(1); i < sortedRecord.NumRows(); i++ { + isDifferent := false + for _, keyIdx := range keyIndices { + prevValScalar := scalar.MakeScalar(sortedRecord.Column(keyIdx), int(i-1)); currValScalar := scalar.MakeScalar(sortedRecord.Column(keyIdx), int(i)) + if cs, needsRelease := prevValScalar.(interface{ Release() }); needsRelease { cs.Release() } + if cs, needsRelease := currValScalar.(interface{ Release() }); needsRelease { cs.Release() } + if !scalar.Equals(prevValScalar, currValScalar) { isDifferent = true; break } + } + if isDifferent { uniqueRowIndices = append(uniqueRowIndices, i) } + } + indicesBuilder := array.NewInt64Builder(adf.mem); defer indicesBuilder.Release(); indicesBuilder.AppendValues(uniqueRowIndices, nil) + indicesArr := indicesBuilder.NewArray(); defer indicesArr.Release() + ctx := compute.WithAllocator(context.Background(), adf.mem) + distinctRecordDatum, err := compute.Take(ctx, compute.TakeOptions{}, arrow.NewRecordDatum(sortedRecord), arrow.NewArrayDatum(indicesArr)) + if err != nil { panic(fmt.Sprintf("Distinct: Take failed: %v", err)) }; defer distinctRecordDatum.Release() + distinctRecord, okValue := distinctRecordDatum.(*arrow.RecordDatum).Value().(arrow.Record); if !okValue { panic("Distinct: Take bad return") } + return NewArrowDataFrameWithAllocator(adf.name, distinctRecord, adf.schema, adf.mem) +} +func (adf *arrowDataFrame) Append(otherRaw df.DataFrame) df.DataFrame { + if otherRaw == nil { panic("Append: other df nil") }; otherArrowDf, ok := otherRaw.(*arrowDataFrame); if !ok { panic(fmt.Sprintf("Append: expected *arrowDataFrame, got %T", otherRaw)) } + currentIsColEmpty := adf.record == nil || adf.record.NumCols() == 0; otherIsColEmpty := otherArrowDf.record == nil || otherArrowDf.record.NumCols() == 0 + currentSchemaForEmpty := adf.schema.schema; if currentSchemaForEmpty == nil || currentSchemaForEmpty.NumFields() != 0 { currentSchemaForEmpty = arrow.NewSchema([]arrow.Field{}, nil) } + if currentIsColEmpty { + if otherIsColEmpty { numRows := adf.Len() + otherArrowDf.Len(); newRec := array.NewRecord(currentSchemaForEmpty, nil, numRows); defer newRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, newRec, adf.schema, adf.mem) } + newOtherRec := otherArrowDf.record.NewSlice(0, otherArrowDf.record.NumRows()); defer newOtherRec.Release(); return NewArrowDataFrameWithAllocator(otherArrowDf.name, newOtherRec, otherArrowDf.schema, adf.mem) + } + if otherIsColEmpty { newThisRec := adf.record.NewSlice(0, adf.record.NumRows()); defer newThisRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, newThisRec, adf.schema, adf.mem) } + if !adf.schema.Equals(otherArrowDf.schema) { panic(fmt.Sprintf("Append: schema mismatch. Current: %s, Other: %s", adf.schema.schema, otherArrowDf.schema.schema)) } + if adf.record.NumRows() == 0 { newOtherRec := otherArrowDf.record.NewSlice(0, otherArrowDf.record.NumRows()); defer newOtherRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, newOtherRec, otherArrowDf.schema, adf.mem) } + if otherArrowDf.record.NumRows() == 0 { newThisRec := adf.record.NewSlice(0, adf.record.NumRows()); defer newThisRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, newThisRec, adf.schema, adf.mem) } + numCols := int(adf.record.NumCols()); concatenatedCols := make([]arrow.Array, numCols); var err error + for i := 0; i < numCols; i++ { + col1 := adf.record.Column(i); col2 := otherArrowDf.record.Column(i) + concatenatedCols[i], err = array.Concatenate([]arrow.Array{col1, col2}, adf.mem) + if err != nil { for j := 0; j < i; j++ { if concatenatedCols[j] != nil { concatenatedCols[j].Release() } }; panic(fmt.Sprintf("Append: concat col %d ('%s'): %v", i, adf.schema.Get(i).Name, err)) } + } + newNumRows := adf.record.NumRows() + otherArrowDf.record.NumRows() + appendedRecord := array.NewRecord(adf.schema.schema, concatenatedCols, newNumRows) + for _, col := range concatenatedCols { if col != nil { col.Release() } }; defer appendedRecord.Release() + return NewArrowDataFrameWithAllocator(adf.name, appendedRecord, adf.schema, adf.mem) +} +func (adf *arrowDataFrame) Union(otherRaw df.DataFrame) df.DataFrame { + if otherRaw == nil { panic("Union: other df nil") }; appendedDf := adf.Append(otherRaw) + unionDf := appendedDf.Distinct() + if appendedArrowDf, ok := appendedDf.(*arrowDataFrame); ok { appendedArrowDf.Release() } + return unionDf +} +func (adf *arrowDataFrame) WhenNil(fillValues map[string]df.Value) df.DataFrame { + if adf.record == nil { panic("WhenNil on nil record") }; if len(fillValues) == 0 { newRec := adf.record.NewSlice(0, adf.record.NumRows()); defer newRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, newRec, adf.schema, adf.mem) } + newRecordCols := make([]arrow.Array, adf.record.NumCols()); modified := false; ctx := compute.WithAllocator(context.Background(), adf.mem) + for i := 0; i < int(adf.record.NumCols()); i++ { + col := adf.record.Column(i); colName := adf.schema.schema.Field(i).Name + fillVal, colNeedsFilling := fillValues[colName] + if !colNeedsFilling || fillVal == nil { col.Retain(); newRecordCols[i] = col; continue } + modified = true; targetArrowType := col.DataType() + // Using dfValueToArrowScalar from types.go (does not take allocator) + fillScalar, err := dfValueToArrowScalar(fillVal, targetArrowType) + if err != nil { + // Release previously retained/created columns before panicking + for j := 0; j < i; j++ { if newRecordCols[j] != nil {newRecordCols[j].Release()} } + panic(fmt.Sprintf("WhenNil: convert fill for '%s': %v", colName, err)) + } + // dfValueToArrowScalar from types.go does not return retained scalars needing release by caller. + + resultDatum, err := compute.FillNull(ctx, arrow.NewArrayDatum(col), arrow.NewScalarDatum(fillScalar)) + if err != nil { + for j := 0; j < i; j++ { if newRecordCols[j] != nil {newRecordCols[j].Release()} } + panic(fmt.Sprintf("WhenNil: FillNull for '%s': %v", colName, err)) + } + newColArr := resultDatum.Value().(arrow.Array); newColArr.Retain() // Retain for newRecordCols + resultDatum.Release(); + newRecordCols[i] = newColArr + } + if !modified { newRec := adf.record.NewSlice(0, adf.record.NumRows()); defer newRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, newRec, adf.schema, adf.mem) } + finalRecord := array.NewRecord(adf.schema.schema, newRecordCols, adf.record.NumRows()) + for _, col := range newRecordCols { col.Release() }; defer finalRecord.Release() + return NewArrowDataFrameWithAllocator(adf.name, finalRecord, adf.schema, adf.mem) +} +func (adf *arrowDataFrame) When(replaceMap map[string]map[any]df.Value) df.DataFrame { + if adf.record == nil { panic("When on nil record") }; if len(replaceMap) == 0 { newRec := adf.record.NewSlice(0, adf.record.NumRows()); defer newRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, newRec, adf.schema, adf.mem) } + newRecordCols := make([]arrow.Array, adf.record.NumCols()); modified := false; ctx := compute.WithAllocator(context.Background(), adf.mem) + for i := 0; i < int(adf.record.NumCols()); i++ { + originalCol := adf.record.Column(i); colName := adf.schema.schema.Field(i).Name; colType := originalCol.DataType() + valueReplacements, colNeedsUpdate := replaceMap[colName] + if !colNeedsUpdate || len(valueReplacements) == 0 { originalCol.Retain(); newRecordCols[i] = originalCol; continue } + modified = true; b := builder.NewBuilder(adf.mem, colType); defer b.Release() + for r := 0; r < originalCol.Len(); r++ { + currentDfVal := adf.GetValue(r, i) + var goKeyForLookup any + if currentDfVal.IsNil() { goKeyForLookup = nil } else { goKeyForLookup = currentDfVal.Get() } + replacementDfVal, shouldReplace := valueReplacements[goKeyForLookup] + if shouldReplace { + // Using dfValueToArrowScalar from types.go (does not take allocator) + replacementScalar, err := dfValueToArrowScalar(replacementDfVal, colType) + if err != nil { + for j := 0; j < i; j++ { if newRecordCols[j] != nil {newRecordCols[j].Release()} } + panic(fmt.Sprintf("When: convert replacement for key %v, col '%s': %v", goKeyForLookup, colName, err)) + } + // Using appendScalarToBuilder from types.go + errAppend := appendScalarToBuilder(b, replacementScalar, colType) + // dfValueToArrowScalar from types.go does not return retained scalars. + if errAppend != nil { + for j := 0; j < i; j++ { if newRecordCols[j] != nil {newRecordCols[j].Release()} } + panic(fmt.Sprintf("When: append replacement for key %v, col '%s': %v", goKeyForLookup, colName, errAppend)) + } + } else { + originalScalarToAppend := currentDfVal.(*arrowValue).val + // Using appendScalarToBuilder from types.go + if err := appendScalarToBuilder(b, originalScalarToAppend, colType); err != nil { + for j := 0; j < i; j++ { if newRecordCols[j] != nil {newRecordCols[j].Release()} } + panic(fmt.Sprintf("When: copy original for col '%s', row %d: %v", colName, r, err)) + } + } + } + newRecordCols[i] = b.NewArray() + } + if !modified { newRec := adf.record.NewSlice(0, adf.record.NumRows()); defer newRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, newRec, adf.schema, adf.mem) } + finalRecord := array.NewRecord(adf.schema.schema, newRecordCols, adf.record.NumRows()) + for _, col := range newRecordCols { col.Release() }; defer finalRecord.Release() + return NewArrowDataFrameWithAllocator(adf.name, finalRecord, adf.schema, adf.mem) +} +func (adf *arrowDataFrame) Join(outputSchemaGiven df.DataFrameSchema, otherRaw df.DataFrame, jointype df.JoinType, joinColsMap map[string]string, fUser func(r1 df.Row, r2 df.Row) []df.Row) df.DataFrame { + if adf.record == nil && !(jointype == df.JoinRight || jointype == df.JoinOuter || df.JoinType(string(jointype)) == "leftanti" ) { // Assuming JoinLeftAnti is a string const + if jointype == df.JoinEqui || jointype == df.JoinLeft { + outputArrowDFSchema, ok := outputSchemaGiven.(*arrowDataFrameSchema); if !ok {panic(fmt.Sprintf("Join: outputSchemaGiven must be *arrowDataFrameSchema, got %T", outputSchemaGiven))} + emptyOutputRec := array.NewRecord(outputArrowDFSchema.schema, nil, 0); defer emptyOutputRec.Release() + return NewArrowDataFrameWithAllocator(adf.name, emptyOutputRec, outputArrowDFSchema, adf.mem) + } + } + if otherRaw == nil { panic("Join: other dataframe cannot be nil") } + otherArrowDf, ok := otherRaw.(*arrowDataFrame); if !ok { panic(fmt.Sprintf("Join: expected *arrowDataFrame, got %T", otherRaw)) } + + if (otherArrowDf.record == nil || otherArrowDf.Len() == 0) && (jointype == df.JoinEqui || jointype == df.JoinLeft || df.JoinType(string(jointype)) == "leftanti") { + outputArrowDFSchema, ok := outputSchemaGiven.(*arrowDataFrameSchema); if !ok {panic(fmt.Sprintf("Join: outputSchemaGiven must be *arrowDataFrameSchema, got %T", outputSchemaGiven))} + schemaForEmpty := outputArrowDFSchema.schema; if df.JoinType(string(jointype)) == "leftanti" { schemaForEmpty = adf.schema.schema } + emptyOutputRec := array.NewRecord(schemaForEmpty, nil, 0); defer emptyOutputRec.Release() + return NewArrowDataFrameWithAllocator(adf.name, emptyOutputRec, outputArrowDFSchema, adf.mem) + } + if (adf.record == nil || adf.Len() == 0) && (jointype == df.JoinRight || jointype == df.JoinEqui) { + outputArrowDFSchema, ok := outputSchemaGiven.(*arrowDataFrameSchema); if !ok {panic(fmt.Sprintf("Join: outputSchemaGiven must be *arrowDataFrameSchema, got %T", outputSchemaGiven))} + emptyOutputRec := array.NewRecord(outputArrowDFSchema.schema, nil, 0); defer emptyOutputRec.Release() + return NewArrowDataFrameWithAllocator(adf.name, emptyOutputRec, outputArrowDFSchema, adf.mem) + } + // For LeftAnti: if adf is empty, result is empty (matching adf schema) + if (adf.record == nil || adf.Len() == 0) && df.JoinType(string(jointype)) == "leftanti" { + emptyOutputRec := array.NewRecord(adf.schema.schema, nil, 0); defer emptyOutputRec.Release() + return NewArrowDataFrameWithAllocator(adf.name, emptyOutputRec, adf.schema, adf.mem) + } + + + if outputSchemaGiven == nil { panic("Join: outputSchemaGiven cannot be nil") } + outputArrowDFSchema, ok := outputSchemaGiven.(*arrowDataFrameSchema); if !ok { panic(fmt.Sprintf("Join: outputSchemaGiven must be *arrowDataFrameSchema, got %T", outputSchemaGiven)) } + if outputArrowDFSchema.schema == nil { panic("Join: outputSchemaGiven's internal arrow.Schema is nil") } + + // Ensure all semi/anti join types are correctly identified for fUser check and logic path. + // df.JoinLeftAnti was previously `JoinLeftAnti` (local const) or `df.JoinType("leftanti")`. Assuming it's now a proper df.JoinType const. + isSemiOrAntiJoin := (jointype == df.JoinLeftAnti || jointype == df.JoinRightAnti || jointype == df.JoinLeftSemi || jointype == df.JoinRightSemi) + if fUser == nil && !isSemiOrAntiJoin { + // fUser is not allowed for Inner, Left, Right, FullOuter, Cross if it's nil. + // It IS allowed to be nil for Semi/Anti joins. + panic("Join: user function fUser cannot be nil for this join type") + } + + + ctx := compute.WithAllocator(context.Background(), adf.mem) + + // Handle CrossJoin separately as it doesn't use keys from joinColsMap + if jointype == df.JoinCross { + if fUser == nil { + panic("Join: CrossJoin requires an fUser function") + } + // Ensure output schema is provided + if outputArrowDFSchema == nil || outputArrowDFSchema.schema == nil { + panic("Join: CrossJoin requires a valid outputSchemaGiven with an internal arrow.Schema") + } + + outputInternalSchema := outputArrowDFSchema.schema + numOutputCols := outputInternalSchema.NumFields() + finalBuilders := make([]array.Builder, numOutputCols) + for i := 0; i < numOutputCols; i++ { + finalBuilders[i] = builder.NewBuilder(adf.mem, outputInternalSchema.Field(i).Type) + } + defer func() { + for _, b := range finalBuilders { + if b != nil { + b.Release() + } + } + }() + + // Handle empty inputs for CrossJoin + if adf.record == nil || adf.record.NumRows() == 0 || otherArrowDf.record == nil || otherArrowDf.record.NumRows() == 0 { + // Result is an empty dataframe with the output schema + emptyCols := make([]arrow.Array, numOutputCols) + for i, b := range finalBuilders { + emptyCols[i] = b.NewArray() + } + finalRecord := array.NewRecord(outputInternalSchema, emptyCols, 0) + for _, col := range emptyCols { + col.Release() + } + defer finalRecord.Release() + return NewArrowDataFrameWithAllocator(adf.name, finalRecord, outputArrowDFSchema, adf.mem) + } + + // Iterate through all combinations of rows + for lIdx := int64(0); lIdx < adf.record.NumRows(); lIdx++ { + leftRowView, errL := NewArrowRowFromRecord(adf.schema, adf.record, int(lIdx)) + if errL != nil { + panic(fmt.Sprintf("Join: CrossJoin error creating left row view for index %d: %v", lIdx, errL)) + } + + for rIdx := int64(0); rIdx < otherArrowDf.record.NumRows(); rIdx++ { + rightRowView, errR := NewArrowRowFromRecord(otherArrowDf.schema, otherArrowDf.record, int(rIdx)) + if errR != nil { + panic(fmt.Sprintf("Join: CrossJoin error creating right row view for index %d: %v", rIdx, errR)) + } + + outputRows := fUser(leftRowView, rightRowView) + for _, outRow := range outputRows { + if outRow.Len() != numOutputCols { + panic(fmt.Sprintf("Join: CrossJoin fUser returned row with %d columns, expected %d", outRow.Len(), numOutputCols)) + } + for c := 0; c < numOutputCols; c++ { + val := outRow.Get(c) + av, ok_av := val.(*arrowValue) + var scalarToAppend scalar.Scalar + if val.IsNil() || !ok_av { // Handle nil or non-arrowValue from fUser by using Null scalar of target type + scalarToAppend = scalar.NewNullScalar(finalBuilders[c].Type()) + } else { + scalarToAppend = av.val + } + // Using 3-argument appendScalarToBuilder from types.go + errAppend := appendScalarToBuilder(finalBuilders[c], scalarToAppend, finalBuilders[c].Type()) + if errAppend != nil { + panic(fmt.Sprintf("Join: CrossJoin append col %d (name: %s): %v. Scalar: %s, BuilderType: %s", c, outputInternalSchema.Field(c).Name, errAppend, scalarToAppend.DataType().Name(), finalBuilders[c].Type().Name())) + } + } + } + } + } + + finalCols := make([]arrow.Array, numOutputCols) + var finalNumRows int64 + if numOutputCols > 0 && finalBuilders[0] != nil { + finalNumRows = int64(finalBuilders[0].Len()) + } + for i, b := range finalBuilders { + if b == nil { + panic(fmt.Sprintf("Join: CrossJoin nil builder at index %d", i)) + } + finalCols[i] = b.NewArray() + } + finalRecord := array.NewRecord(outputInternalSchema, finalCols, finalNumRows) + for _, col := range finalCols { + col.Release() + } + defer finalRecord.Release() + return NewArrowDataFrameWithAllocator(adf.name, finalRecord, outputArrowDFSchema, adf.mem) + } + + leftKeyDatums := make([]arrow.Datum, 0, len(joinColsMap)) + rightKeyDatums := make([]arrow.Datum, 0, len(joinColsMap)) + if (adf.record == nil || otherArrowDf.record == nil) && len(joinColsMap) > 0 { + panic("Join: Cannot prepare keys for join as one or both records are nil") + } + + for lKeyName, rKeyName := range joinColsMap { + lIdx := adf.schema.GetIndexByName(lKeyName); if lIdx == -1 { panic(fmt.Sprintf("Join: left key '%s' not found", lKeyName)) } + leftKeyDatums = append(leftKeyDatums, arrow.NewArrayDatum(adf.record.Column(lIdx))) + rIdx := otherArrowDf.schema.GetIndexByName(rKeyName); if rIdx == -1 { panic(fmt.Sprintf("Join: right key '%s' not found", rKeyName)) } + rightKeyDatums = append(rightKeyDatums, arrow.NewArrayDatum(otherArrowDf.record.Column(rIdx))) + if !arrow.TypeEqual(adf.record.Column(lIdx).DataType(), otherArrowDf.record.Column(rIdx).DataType()) { + panic(fmt.Sprintf("Join: type mismatch for key. L:'%s'(%s) R:'%s'(%s)", lKeyName, adf.record.Column(lIdx).DataType(), rKeyName, otherArrowDf.record.Column(rIdx).DataType())) + } + } + if len(leftKeyDatums) == 0 { panic("JoinEqui and other key-based joins require join columns.") } + defer func() { for _,d := range leftKeyDatums { d.Release() }; for _,d := range rightKeyDatums {d.Release()} }() + + var hjComputeJoinType compute.JoinType + switch jointype { + case df.JoinEqui: hjComputeJoinType = compute.InnerJoin + case df.JoinLeft: hjComputeJoinType = compute.LeftOuterJoin + case df.JoinRight: hjComputeJoinType = compute.RightOuterJoin + case df.JoinOuter: hjComputeJoinType = compute.FullOuterJoin + // Assuming df.JoinLeftAnti, etc., are defined constants of type df.JoinType + case df.JoinLeftAnti: hjComputeJoinType = compute.LeftAntiJoin + case df.JoinLeftSemi: hjComputeJoinType = compute.LeftSemiJoin + case df.JoinRightSemi: hjComputeJoinType = compute.RightSemiJoin + case df.JoinRightAnti: hjComputeJoinType = compute.RightAntiJoin + default: + // If fUser is nil here, it means it was a semi/anti join not caught above, which is an issue. + // Or, it's a non-semi/anti join type that's not supported by HashJoin path. + if fUser == nil && isSemiOrAntiJoin { // Should have been caught by the switch + panic(fmt.Sprintf("Join: internal error - semi/anti join type %s not mapped for HashJoin", jointype)) + } + panic(fmt.Sprintf("Join: unsupported join type %s for HashJoin path", jointype)) + } + + hjIndicesTable, err := compute.HashJoin(ctx, leftKeyDatums, rightKeyDatums, + arrow.NewRecordDatum(adf.record), arrow.NewRecordDatum(otherArrowDf.record), + hjComputeJoinType, compute.HashJoinOptions{LeftSuffix:"_L", RightSuffix:"_R"}) + if err != nil { panic(fmt.Sprintf("Join: HashJoin compute failed for type %s: %v", jointype, err)) } + defer hjIndicesTable.Release() + + if isSemiOrAntiJoin { + if hjIndicesTable.NumCols() != 1 { panic(fmt.Sprintf("Join: %s HashJoin result expected 1 col (indices), got %d", jointype, hjIndicesTable.NumCols())) } + + var sourceRecordForTake arrow.Record + var sourceSchemaForOutput *arrowDataFrameSchema + + // Determine which table's rows to output based on the join type + if jointype == df.JoinLeftSemi || jointype == df.JoinLeftAnti { + sourceRecordForTake = adf.record + sourceSchemaForOutput = adf.schema + } else if jointype == df.JoinRightSemi || jointype == df.JoinRightAnti { + sourceRecordForTake = otherArrowDf.record + sourceSchemaForOutput = otherArrowDf.schema + } else { + panic(fmt.Sprintf("Join: internal error - unhandled semi/anti join type %s in output determination logic", jointype)) + } + + // If the designated source table for the semi/anti join is nil or empty, the result is also empty, + // but with the schema of that source table. + if sourceRecordForTake == nil || sourceRecordForTake.NumRows() == 0 { + emptyFinalRecord := array.NewRecord(sourceSchemaForOutput.schema, nil, 0) + // No defer release for emptyFinalRecord if it's immediately returned and not retained elsewhere. + // NewArrowDataFrameWithAllocator will handle it. + return NewArrowDataFrameWithAllocator(adf.name, emptyFinalRecord, sourceSchemaForOutput, adf.mem) + } + + hjTr, errTr := array.NewTableReader(hjIndicesTable, -1); + if errTr != nil { panic(fmt.Sprintf("Join: Failed to create TableReader for HashJoin result for %s: %v", jointype, errTr)) } + defer hjTr.Release() + + var finalRecord arrow.Record + + if hjTr.Next() { // Check if there's at least one batch of indices + indicesRecord := hjTr.Record() // This record contains a single column of indices. + indicesArr := indicesRecord.Column(0) + + if indicesArr.Len() > 0 { + takenDatum, errTake := compute.Take(ctx, compute.TakeOptions{}, + arrow.NewRecordDatum(sourceRecordForTake), + arrow.NewArrayDatum(indicesArr)) + if errTake != nil { panic(fmt.Sprintf("Join: %s Take failed: %v", jointype, errTake)) } + + resultValue := takenDatum.Value() + if resultValue == nil { takenDatum.Release(); panic(fmt.Sprintf("Join: %s Take result datum value is nil", jointype)) } + recResult, okRec := resultValue.(arrow.Record) + if !okRec { takenDatum.Release(); panic(fmt.Sprintf("Join: %s Take did not return arrow.Record, got %T", jointype, resultValue)) } + // recResult is effectively owned by takenDatum. We need our own ref if takenDatum is released. + recResult.Retain() + takenDatum.Release() + finalRecord = recResult // Now we own this reference + } else { + // No indices found by HashJoin (e.g., no matches), result is empty. + finalRecord = array.NewRecord(sourceSchemaForOutput.schema, nil, 0) + } + } else { // No records in hjIndicesTable + if hjTr.Err() != nil { panic(fmt.Sprintf("Join: error reading %s HashJoin indices: %v", jointype, hjTr.Err())) } + finalRecord = array.NewRecord(sourceSchemaForOutput.schema, nil, 0) + } + // finalRecord is now either a record with data (retained) or an empty record (newly created). + // NewArrowDataFrameWithAllocator will retain it again. So, we must release our hold here. + defer finalRecord.Release() + return NewArrowDataFrameWithAllocator(adf.name, finalRecord, sourceSchemaForOutput, adf.mem) + } + + // Path for INNER, LEFT, RIGHT, FULL OUTER (uses fUser) + if hjIndicesTable.NumCols() != 2 { panic(fmt.Sprintf("Join: HashJoin result expected 2 index cols for %s, got %d", jointype, hjIndicesTable.NumCols()))} + outputInternalSchema := outputArrowDFSchema.schema; numOutputCols := outputInternalSchema.NumFields() + finalBuilders := make([]array.Builder, numOutputCols); + for i:=0; i 0 && finalBuilders[0] != nil { finalNumRows = int64(finalBuilders[0].Len()) } + for i, b := range finalBuilders { if b == nil {panic(fmt.Sprintf("Join: nil builder at index %d",i))}; finalCols[i] = b.NewArray() } + finalRecord := array.NewRecord(outputInternalSchema, finalCols, finalNumRows) + for _, col := range finalCols { col.Release() }; defer finalRecord.Release() + return NewArrowDataFrameWithAllocator(adf.name, finalRecord, outputArrowDFSchema, adf.mem) +} + +func (adf *arrowDataFrame) Intersection(otherRaw df.DataFrame, cols ...string) df.DataFrame { + if adf.record == nil { schemaToUse := adf.schema; if adf.schema == nil || adf.schema.schema == nil { emptyFields := []arrow.Field{}; schemaToUse = NewArrowDataFrameSchema(arrow.NewSchema(emptyFields,nil)).(*arrowDataFrameSchema) }; emptyRec := array.NewRecord(schemaToUse.schema, nil, 0); defer emptyRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, emptyRec, schemaToUse, adf.mem) } + if otherRaw == nil { panic("Intersection: other dataframe cannot be nil") } + otherArrowDf, ok := otherRaw.(*arrowDataFrame); if !ok { panic(fmt.Sprintf("Intersection: expected *arrowDataFrame, got %T", otherRaw)) } + if otherArrowDf.record == nil { emptyRec := array.NewRecord(adf.schema.schema, nil, 0); defer emptyRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, emptyRec, adf.schema, adf.mem) } + actualJoinColsMap := make(map[string]string) + if len(cols) == 0 { + commonColsFound := false + for _, name1 := range adf.schema.Names() { + idx2 := otherArrowDf.schema.GetIndexByName(name1) + if idx2 != -1 { + type1 := adf.schema.schema.Field(adf.schema.GetIndexByName(name1)).Type; type2 := otherArrowDf.schema.schema.Field(idx2).Type + if arrow.TypeEqual(type1, type2) { actualJoinColsMap[name1] = name1; commonColsFound = true } + } + } + if !commonColsFound { emptyRec := array.NewRecord(adf.schema.schema, nil, 0); defer emptyRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, emptyRec, adf.schema, adf.mem) } + } else { + for _, colName := range cols { + if !adf.schema.HasName(colName) { panic(fmt.Sprintf("Intersection: key col '%s' not in left df", colName)) } + if !otherArrowDf.schema.HasName(colName) { panic(fmt.Sprintf("Intersection: key col '%s' not in right df", colName)) } + type1 := adf.schema.schema.Field(adf.schema.GetIndexByName(colName)).Type; type2 := otherArrowDf.schema.schema.Field(otherArrowDf.schema.GetIndexByName(colName)).Type + if !arrow.TypeEqual(type1, type2) { panic(fmt.Sprintf("Intersection: type mismatch for key col '%s'. L: %s, R: %s", colName, type1, type2)) } + actualJoinColsMap[colName] = colName + } + } + if len(actualJoinColsMap) == 0 { emptyRec := array.NewRecord(adf.schema.schema, nil, 0); defer emptyRec.Release(); return NewArrowDataFrameWithAllocator(adf.name, emptyRec, adf.schema, adf.mem) } + fSelectLeftRow := func(r1 df.Row, r2 df.Row) []df.Row { return []df.Row{r1} } + joinedDf := adf.Join(adf.schema, otherArrowDf, df.JoinEqui, actualJoinColsMap, fSelectLeftRow) + distinctResultDf := joinedDf.Distinct() + if arrowJoinedDf, ok_join := joinedDf.(*arrowDataFrame); ok_join { arrowJoinedDf.Release() } + return distinctResultDf +} +func (adf *arrowDataFrame) GroupBy(cols ...string) df.GroupedDataFrame { + if adf.record == nil { panic("GroupBy called on nil dataframe record") } + if len(cols) == 0 { panic("GroupBy requires at least one column name") } + keyIndices := make([]int, len(cols)) + for i, name := range cols { + idx := adf.schema.GetIndexByName(name) + if idx == -1 { panic(fmt.Sprintf("GroupBy: column '%s' not found", name)) } + keyIndices[i] = idx + } + ctx := compute.WithAllocator(context.Background(), adf.mem) + keyColsFromRec := make([]arrow.Array, len(keyIndices)); keyFieldsForSchema := make([]arrow.Field, len(keyIndices)) + for i, ki := range keyIndices { keyColsFromRec[i] = adf.record.Column(ki); keyColsFromRec[i].Retain(); keyFieldsForSchema[i] = adf.schema.schema.Field(ki) } + keysOnlyRecSchema := arrow.NewSchema(keyFieldsForSchema, nil) + keysOnlyRec := array.NewRecord(keysOnlyRecSchema, keyColsFromRec, adf.record.NumRows()) + for _, col := range keyColsFromRec { col.Release() }; defer keysOnlyRec.Release() + tblReader, err := array.NewRecordReader(keysOnlyRecSchema, []arrow.Record{keysOnlyRec}) + if err != nil { panic(fmt.Sprintf("GroupBy: failed to create record reader for keys: %v", err)) }; defer tblReader.Release() + keysOnlyTable, err := array.NewTableFromReader(tblReader, -1) + if err != nil { panic(fmt.Sprintf("GroupBy: failed to create keysOnlyTable: %v", err)) }; defer keysOnlyTable.Release() + keyColNames := make([]string, len(keyIndices)); for i, ki := range keyIndices { keyColNames[i] = adf.schema.schema.Field(ki).Name } + uniqueKeysResultTable, err := keysOnlyTable.Distinct(ctx, keyColNames...) + if err != nil { panic(fmt.Sprintf("GroupBy: failed to get distinct keys: %v", err)) }; + adf.record.Retain() + return &arrowGroupedDataFrame{ originalRecord: adf.record, originalSchema: adf.schema, groupingColNames: cols, uniqueKeysTable: uniqueKeysResultTable, mem: adf.mem, } +} + +func (adf *arrowDataFrame) Except(otherRaw df.DataFrame, cols ...string) df.DataFrame { + if adf.record == nil || adf.record.NumRows() == 0 { + schemaToUse := adf.schema + if adf.schema == nil || adf.schema.schema == nil { + emptyFields := []arrow.Field{} + schemaToUse = NewArrowDataFrameSchema(arrow.NewSchema(emptyFields,nil)).(*arrowDataFrameSchema) + } + emptyRec := array.NewRecord(schemaToUse.schema, nil, 0); defer emptyRec.Release() + return NewArrowDataFrameWithAllocator(adf.name, emptyRec, schemaToUse, adf.mem) + } + if otherRaw == nil { panic("Except: other dataframe cannot be nil") } + otherArrowDf, ok := otherRaw.(*arrowDataFrame) + if !ok { panic(fmt.Sprintf("Except: expected *arrowDataFrame, got %T", otherRaw)) } + if otherArrowDf.record == nil || otherArrowDf.record.NumRows() == 0 { return adf.Distinct() } + + actualJoinColsMap := make(map[string]string) + if len(cols) == 0 { + commonColsFound := false + for _, name1 := range adf.schema.Names() { + idx2 := otherArrowDf.schema.GetIndexByName(name1) + if idx2 != -1 { + type1 := adf.schema.schema.Field(adf.schema.GetIndexByName(name1)).Type + type2 := otherArrowDf.schema.schema.Field(idx2).Type + if arrow.TypeEqual(type1, type2) { actualJoinColsMap[name1] = name1; commonColsFound = true } + } + } + // If no common columns with same type found for implicit join, result is all distinct rows of left. + if !commonColsFound { return adf.Distinct() } + } else { + for _, colName := range cols { + if !adf.schema.HasName(colName) { panic(fmt.Sprintf("Except: key col '%s' not in left df", colName)) } + if !otherArrowDf.schema.HasName(colName) { panic(fmt.Sprintf("Except: key col '%s' not in right df", colName)) } + type1 := adf.schema.schema.Field(adf.schema.GetIndexByName(colName)).Type + type2 := otherArrowDf.schema.schema.Field(otherArrowDf.schema.GetIndexByName(colName)).Type + if !arrow.TypeEqual(type1, type2) { panic(fmt.Sprintf("Except: type mismatch for key '%s'", colName)) } + actualJoinColsMap[colName] = colName + } + } + // If no join columns specified AND no common columns found (for implicit all-column join), + // or if explicit cols were empty, it implies an except based on full row comparison. + // The current HashJoin path requires specific key columns. + // If actualJoinColsMap is empty but cols was also empty (meaning all-cols implicit join), + // this means we need to use all columns as keys if they are compatible. + // This part of logic for "all columns" except might need more specific handling if actualJoinColsMap remains empty. + // For now, if no join keys, and it's not a zero-column DF, result is adf.Distinct(). + if len(actualJoinColsMap) == 0 && adf.record.NumCols() > 0 { return adf.Distinct() } + if adf.record.NumCols() == 0 { return adf.Distinct() } // Except on an empty-column DF is itself distinct. + + // Use Join with JoinLeftAnti. fUser is nil as LeftAntiJoin produces rows from left table. + // Output schema is the left table's schema. + // Assuming JoinLeftAnti is defined in the df package or as a recognized string const by Join. + leftAntiJoinedDf := adf.Join(adf.schema, otherArrowDf, df.JoinType("leftanti"), actualJoinColsMap, nil) + + // The result of LeftAntiJoin already contains rows from 'adf' not in 'other'. + // Now, make these rows distinct. + resultDf := leftAntiJoinedDf.Distinct() + + if arrowJoinedDf, ok_join := leftAntiJoinedDf.(*arrowDataFrame); ok_join { + arrowJoinedDf.Release() + } + return resultDf +} + +// evaluateExpr is a conceptual helper. The logic will be inlined into Select or a private method. +// func (adf *arrowDataFrame) evaluateExpr(expr df.Expr) (arrow.Array, df.SeriesSchema, error) { ... } + +func (adf *arrowDataFrame) Select(expressions ...df.Expr) df.DataFrame { + if adf.record == nil && len(expressions) > 0 { + panic("Select: cannot select from a DataFrame with a nil record") + } + + numRows := int64(0) + if adf.record != nil { + numRows = adf.record.NumRows() + } + + if len(expressions) == 0 { // Return a DataFrame with 0 columns but same number of rows + emptyArrowSchema := arrow.NewSchema([]arrow.Field{}, nil) + emptyDfSchema := NewArrowDataFrameSchema(emptyArrowSchema).(*arrowDataFrameSchema) + emptyRecord := array.NewRecord(emptyArrowSchema, nil, numRows) + defer emptyRecord.Release() + return NewArrowDataFrameWithAllocator(adf.name, emptyRecord, emptyDfSchema, adf.mem) + } + + newArrays := make([]arrow.Array, len(expressions)) + newFields := make([]arrow.Field, len(expressions)) + ctx := compute.WithAllocator(context.Background(), adf.mem) + + + // Helper to clean up arrays created so far in case of an error + cleanupArraysOnError := func(count int) { + for k := 0; k < count; k++ { + if newArrays[k] != nil { + newArrays[k].Release() + } + } + } + + for i, expr := range expressions { + outputColName := expr.Name() + var currentResultArray arrow.Array + var currentResultSeriesSchema df.SeriesSchema + + switch expr.OpType() { + case df.ColNameExpr: + colName := expr.Col() + if colName == "" { cleanupArraysOnError(i); panic(fmt.Sprintf("Select: column name expression for expr %d ('%s') is empty", i, outputColName)) } + originalSeries := adf.GetSeriesByName(colName) + arrowS, ok := originalSeries.(*arrowSeries) + if !ok { cleanupArraysOnError(i); panic(fmt.Sprintf("Select: expected *arrowSeries for col '%s', got %T", colName, originalSeries)) } + arrowS.arr.Retain(); currentResultArray = arrowS.arr + currentResultSeriesSchema = arrowS.schema + if outputColName == "" { outputColName = colName } // Default to original name if no alias + + case df.LiteralExpr: + literalValue := expr.Const() + if literalValue == nil { cleanupArraysOnError(i); panic(fmt.Sprintf("Select: literal expression for expr %d ('%s') has nil df.Value", i, outputColName)) } + arrowType, err := dfFormatToArrowType(literalValue.Schema().Format) + if err != nil { cleanupArraysOnError(i); panic(fmt.Sprintf("Select: error converting literal format %v to Arrow type for expr %d ('%s'): %v", literalValue.Schema().Format, i, outputColName, err)) } + builder := array.NewBuilder(adf.mem, arrowType) + litScalar, errScalar := dfValueToArrowScalar(literalValue, arrowType) + if errScalar != nil { builder.Release(); cleanupArraysOnError(i); panic(fmt.Sprintf("Select: failed to convert literal value to Arrow scalar for expr %d ('%s'): %v", i, outputColName, errScalar)) } + for r := int64(0); r < numRows; r++ { + errAppend := appendScalarToBuilder(builder, litScalar, arrowType) + if errAppend != nil { builder.Release(); cleanupArraysOnError(i); panic(fmt.Sprintf("Select: error appending literal scalar for expr %d ('%s'): %v", i, outputColName, errAppend)) } + } + currentResultArray = builder.NewArray(); builder.Release() + if outputColName == "" { outputColName = fmt.Sprintf("_literal_%d", i) } + currentResultSeriesSchema = df.SeriesSchema{Name: outputColName, Format: literalValue.Schema().Format, Nullable: literalValue.IsNil()} + + default: // Potentially an operation on a parent column or between columns + if expr.Parent() != nil && expr.Parent().OpType() == df.ColNameExpr { + // This is a unary operation on a column, e.g., Col("A").SomeOp() + parentColName := expr.Parent().Col() + parentSeries := adf.GetSeriesByName(parentColName).(*arrowSeries) // Panics if not found or not arrowSeries + + // Delegate to series.Select. The current 'expr' is the operation node. + // Example: expr = OpConst_Add(5), expr.Parent() = Col("A") + // We call parentSeries(ColA).Select(OpConst_Add(5)) + resultSeries := parentSeries.Select(expr) // This will execute the logic in series.Select + defer resultSeries.Release() + + arrowResultSeries, ok := resultSeries.(*arrowSeries) + if !ok { cleanupArraysOnError(i); parentSeries.Release(); panic("Series.Select did not return *arrowSeries") } + + arrowResultSeries.arr.Retain(); currentResultArray = arrowResultSeries.arr + currentResultSeriesSchema = arrowResultSeries.schema + if outputColName == "" { outputColName = currentResultSeriesSchema.Name } // Use name from series op if not aliased higher up + parentSeries.Release() + + } else if expr.MapOp() != nil && len(expr.MapOp().Args()) > 0 && expr.MapOp().Args()[0].OpType() == df.ColNameExpr { + // Binary operation between a parent column (or the df context if no parent) and another column + // Example: Col("A").Op_Add(Col("B")) -> expr is Op_Add, Parent is Col("A"), MapOp.Arg[0] is Col("B") + + leftSeriesParent := expr.Parent() + if leftSeriesParent == nil || leftSeriesParent.OpType() != df.ColNameExpr { + cleanupArraysOnError(i); panic(fmt.Sprintf("Select: binary op expected parent ColNameExpr for expr '%s'", expr.Name())) + } + leftColName := leftSeriesParent.Col() + leftSeries := adf.GetSeriesByName(leftColName).(*arrowSeries); defer leftSeries.Release() + + rightColName := expr.MapOp().Args()[0].Col() + rightSeries := adf.GetSeriesByName(rightColName).(*arrowSeries); defer rightSeries.Release() + + if leftSeries.Len() != rightSeries.Len() { + cleanupArraysOnError(i); panic(fmt.Sprintf("Select: column length mismatch for binary op between '%s' and '%s'", leftColName, rightColName)) + } + + leftDatum := arrow.NewArrayDatum(leftSeries.arr); defer leftDatum.Release() + rightDatum := arrow.NewArrayDatum(rightSeries.arr); defer rightDatum.Release() + + var computeErr error; var outputDatum arrow.Datum + // Assume expr.Name() or mapOp.Name() gives the binary operation like "Op_Add", "Op_Multiply" etc. + // This part needs to align with how binary ops between columns are defined in df.Expr + opName := expr.Name() // Or from mapOp if that's where the direct op name is stored. + + switch opName { + case "Op_Add": // Hypothetical name for Col("A").Add(Col("B")) + outputDatum, computeErr = compute.Add(ctx, leftDatum, rightDatum, compute.ArithmeticOptions{NoSignedOverflow: false}) + case "Op_Subtract": + outputDatum, computeErr = compute.Subtract(ctx, leftDatum, rightDatum, compute.ArithmeticOptions{NoSignedOverflow: false}) + case "Op_Multiply": + outputDatum, computeErr = compute.Multiply(ctx, leftDatum, rightDatum, compute.ArithmeticOptions{NoSignedOverflow: false}) + case "Op_Divide": + outputDatum, computeErr = compute.Divide(ctx, leftDatum, rightDatum, compute.ArithmeticOptions{NoSignedOverflow: false}) + // Add comparison ops if they follow this pattern too e.g. Col("A").Eq(Col("B")) + case "Op_Eq": // Hypothetical for Col("A").Eq(Col("B")) + outputDatum, computeErr = compute.Compare(ctx,leftDatum,rightDatum,compute.CompareOptions{Operator: compute.EQUAL}) + // ... other binary ops ... + default: + cleanupArraysOnError(i); panic(fmt.Sprintf("Select: unsupported binary operation '%s' between columns for expr '%s'", opName, expr.Name())) + } + + if computeErr != nil { cleanupArraysOnError(i); panic(fmt.Sprintf("Select: compute error for binary op '%s' expr '%s': %v", opName, expr.Name(), computeErr)) } + defer outputDatum.Release() + currentResultArray = outputDatum.MakeArray() + + // Schema for binary op result: Name from expr, Format from result array, Nullable from result array + // Type promotion (e.g. int + float = float) is handled by Arrow compute kernel. + resultArrowType := currentResultArray.DataType() + resultFormat := arrowToDfFormat(resultArrowType) + if outputColName == "" { outputColName = fmt.Sprintf("_result_%d", i) } + currentResultSeriesSchema = df.SeriesSchema{Name: outputColName, Format: resultFormat, Nullable: currentResultArray.NullN() > 0} + + } else { + cleanupArraysOnError(i) + panic(fmt.Sprintf("Select: unsupported complex expression type or structure for expr %d ('%s')", i, outputColName)) + } + } + newArrays[i] = currentResultArray // Already retained + newFields[i] = arrow.Field{Name: outputColName, Type: currentResultArray.DataType(), Nullable: currentResultSeriesSchema.Nullable, Metadata: arrow.MetadataFrom(currentResultSeriesSchema.Metadata)} + } + + defer func() { for _, arr := range newArrays { if arr != nil { arr.Release() } } }() + + finalArrowSchema := arrow.NewSchema(newFields, adf.schema.schema.Metadata()) + finalDfSchema := NewArrowDataFrameSchema(finalArrowSchema).(*arrowDataFrameSchema) + finalRecord := array.NewRecord(finalArrowSchema, newArrays, numRows) + defer finalRecord.Release() + + return NewArrowDataFrameWithAllocator(adf.name, finalRecord, finalDfSchema, adf.mem) +} + + +func (adf *arrowDataFrame) Rename(name string, inplace bool) df.DataFrame { + if name == "" { + panic("DataFrame name cannot be empty") + } + if inplace { + adf.name = name + return adf + } + + if adf.record == nil { + return NewArrowDataFrameWithAllocator(name, nil, adf.schema, adf.mem) + } + return NewArrowDataFrameWithAllocator(name, adf.record, adf.schema, adf.mem) +} + +func (adf *arrowDataFrame) AsFormat(targetFormats map[string]df.Format) df.DataFrame { + if adf.record == nil { + panic("AsFormat: cannot operate on DataFrame with a nil record") + } + if len(targetFormats) == 0 { + return NewArrowDataFrameWithAllocator(adf.name, adf.record, adf.schema, adf.mem) + } + + numCols := int(adf.record.NumCols()) + newArrays := make([]arrow.Array, numCols) + newFields := make([]arrow.Field, numCols) + copy(newFields, adf.schema.schema.Fields()) + + changed := false + ctx := compute.WithAllocator(context.Background(), adf.mem) + + cleanupArraysOnError := func(count int) { + for k := 0; k < count; k++ { if newArrays[k] != nil { newArrays[k].Release() } } + } + + for i := 0; i < numCols; i++ { + originalCol := adf.record.Column(i) + originalField := adf.schema.schema.Field(i) + colName := originalField.Name + + targetFormat, ok := targetFormats[colName] + if !ok || targetFormat == nil { + originalCol.Retain(); newArrays[i] = originalCol; continue + } + currentFormat := adf.schema.Get(i).Format + if currentFormat == targetFormat { + originalCol.Retain(); newArrays[i] = originalCol; continue + } + + changed = true + targetArrowType, err := dfFormatToArrowType(targetFormat) + if err != nil { + cleanupArraysOnError(i) + panic(fmt.Sprintf("AsFormat: error converting target format %v for column '%s' to Arrow type: %v", targetFormat, colName, err)) + } + + if arrow.TypeEqual(originalCol.DataType(), targetArrowType) { + originalCol.Retain(); newArrays[i] = originalCol + } else { + castOptions := compute.DefaultCastOptions(false) + castedArray, castErr := compute.Cast(ctx, originalCol, targetArrowType, castOptions) + if castErr != nil { + cleanupArraysOnError(i) + panic(fmt.Sprintf("AsFormat: error casting column '%s' from %s to %s: %v", colName, originalCol.DataType(), targetArrowType, castErr)) + } + newArrays[i] = castedArray + } + newFields[i].Type = targetArrowType + newFields[i].Nullable = newArrays[i].NullN() > 0 || originalField.Nullable + } + + if !changed { + for _, arr := range newArrays { if arr != nil { arr.Release() } } // Release retained original arrays + return NewArrowDataFrameWithAllocator(adf.name, adf.record, adf.schema, adf.mem) + } + + defer func() { for _, arr := range newArrays { if arr != nil { arr.Release() } } }() + + newArrowSchemaInternal := arrow.NewSchema(newFields, adf.schema.schema.Metadata()) + newDfSchema := NewArrowDataFrameSchema(newArrowSchemaInternal).(*arrowDataFrameSchema) + + newRecord := array.NewRecord(newArrowSchemaInternal, newArrays, adf.record.NumRows()) + defer newRecord.Release() + + return NewArrowDataFrameWithAllocator(adf.name, newRecord, newDfSchema, adf.mem) +} + +func (adf *arrowDataFrame) UpdateSeries(index int, series df.Series) df.DataFrame { + if adf.record == nil { panic("UpdateSeries: cannot update series on a DataFrame with a nil record") } + if index < 0 || index >= int(adf.record.NumCols()) { panic(fmt.Sprintf("UpdateSeries: index %d out of bounds for %d columns", index, adf.record.NumCols())) } + if series == nil { panic("UpdateSeries: input series cannot be nil") } + arrowSeriesToUpdate, ok := series.(*arrowSeries) + if !ok { panic(fmt.Sprintf("UpdateSeries: expected *arrowSeries, got %T", series)) } + if arrowSeriesToUpdate.arr == nil { panic("UpdateSeries: input arrowSeries has a nil internal array") } + if arrowSeriesToUpdate.Len() != adf.Len() { panic(fmt.Sprintf("UpdateSeries: length mismatch. DataFrame has %d rows, input series has %d rows", adf.Len(), arrowSeriesToUpdate.Len())) } + + newRecordCols := make([]arrow.Array, adf.record.NumCols()) + for i, col := range adf.record.Columns() { + if i == index { arrowSeriesToUpdate.arr.Retain(); newRecordCols[i] = arrowSeriesToUpdate.arr + } else { col.Retain(); newRecordCols[i] = col } + } + defer func() { for _, col := range newRecordCols { col.Release() } }() + + newSchemaFields := make([]arrow.Field, adf.record.NumCols()) + copy(newSchemaFields, adf.schema.schema.Fields()) + inputSeriesSchema := arrowSeriesToUpdate.Schema() + inputArrowType, err := dfFormatToArrowType(inputSeriesSchema.Format) + if err != nil { panic(fmt.Sprintf("UpdateSeries: could not convert input series format %v to arrow type: %v", inputSeriesSchema.Format, err)) } + newSchemaFields[index] = arrow.Field{ Name: inputSeriesSchema.Name, Type: inputArrowType, Nullable: inputSeriesSchema.Nullable, Metadata: arrow.MetadataFrom(inputSeriesSchema.Metadata)} + for i, field := range newSchemaFields { if i != index && field.Name == inputSeriesSchema.Name { panic(fmt.Sprintf("UpdateSeries: new series name '%s' conflicts with existing column at index %d", inputSeriesSchema.Name, i)) } } + + newArrowSchema := arrow.NewSchema(newSchemaFields, adf.schema.schema.Metadata()) + newDfSchema := NewArrowDataFrameSchema(newArrowSchema).(*arrowDataFrameSchema) + newRecord := array.NewRecord(newArrowSchema, newRecordCols, adf.record.NumRows()) + defer newRecord.Release() + return NewArrowDataFrameWithAllocator(adf.name, newRecord, newDfSchema, adf.mem) +} + +func (adf *arrowDataFrame) UpdateSeriesByName(name string, series df.Series) df.DataFrame { + idx := adf.schema.GetIndexByName(name) + if idx == -1 { panic(fmt.Sprintf("UpdateSeriesByName: column '%s' not found", name)) } + return adf.UpdateSeries(idx, series) +} + +func (adf *arrowDataFrame) ForEachRow(f func(df.Row)) { + if f == nil { panic("ForEachRow: function f cannot be nil") } + if adf.record == nil || adf.record.NumRows() == 0 { return } + for i := int64(0); i < adf.record.NumRows(); i++ { + rowView, err := NewArrowRowFromRecord(adf.schema, adf.record, int(i)) + if err != nil { panic(fmt.Sprintf("ForEachRow: error creating row view for index %d: %v", i, err)) } + f(rowView) + } +} + +var _ df.DataFrame = (*arrowDataFrame)(nil) diff --git a/df/arrow/df_benchmark_test.go b/df/arrow/df_benchmark_test.go new file mode 100644 index 0000000..259d1b0 --- /dev/null +++ b/df/arrow/df_benchmark_test.go @@ -0,0 +1,307 @@ +//go:build arrow + +package arrow_test + +import ( + "fmt" + "math/rand" + "testing" + "time" + + "github.com/apache/arrow/go/v14/arrow" + "github.com/apache/arrow/go/v14/arrow/array" + "github.com/apache/arrow/go/v14/arrow/memory" + "github.com/apache/arrow/go/v14/arrow/scalar" // For creating Arrow values if needed in fUser etc. + + "github.com/blue4209211/pq/df" + arrowimpl "github.com/blue4209211/pq/df/arrow" + inmemory "github.com/blue4209211/pq/df/inmemory" +) + +const ( + benchmarkNumRows = 100_000 // Default number of rows for benchmarks + benchmarkNumGroups = 100 // Number of distinct groups for cat_string + benchmarkJoinNumRowsL = 100_000 + benchmarkJoinNumRowsR = 50_000 +) + +var volatileDF df.DataFrame // To prevent compiler optimizing out benchmarked operations + +// generateArrowData creates an arrowDataFrame for benchmarking. +// Schema: id_int (Int64), val_float (Float64), cat_string (String), key_join (Int64) +func generateArrowData(b *testing.B, numRows int, mem memory.Allocator) df.DataFrame { + b.Helper() // Marks this as a benchmark helper function + + schema := arrow.NewSchema( + []arrow.Field{ + {Name: "id_int", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, + {Name: "val_float", Type: arrow.PrimitiveTypes.Float64, Nullable: false}, + {Name: "cat_string", Type: arrow.BinaryTypes.String, Nullable: false}, + {Name: "key_join", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, + }, + nil, + ) + dfSchema := arrowimpl.NewArrowDataFrameSchema(schema).(*arrowimpl.ArrowDataFrameSchema) + + idBuilder := array.NewInt64Builder(mem) + defer idBuilder.Release() + valBuilder := array.NewFloat64Builder(mem) + defer valBuilder.Release() + catBuilder := array.NewStringBuilder(mem) + defer catBuilder.Release() + keyJoinBuilder := array.NewInt64Builder(mem) + defer keyJoinBuilder.Release() + + randSrc := rand.New(rand.NewSource(time.Now().UnixNano())) + + for i := 0; i < numRows; i++ { + idBuilder.Append(int64(i)) + valBuilder.Append(randSrc.Float64() * 1000) + catBuilder.Append(fmt.Sprintf("group_%d", i%benchmarkNumGroups)) + keyJoinBuilder.Append(int64(randSrc.Intn(numRows / 2))) // Ensure some key overlap for joins + } + + cols := []arrow.Array{ + idBuilder.NewArray(), + valBuilder.NewArray(), + catBuilder.NewArray(), + keyJoinBuilder.NewArray(), + } + // Release arrays after record takes ownership (or if record creation fails) + defer func() { + for _, col := range cols { + col.Release() + } + }() + + record := array.NewRecord(schema, cols, int64(numRows)) + // NewArrowDataFrameWithAllocator will retain the record. We created it, so we must release it. + defer record.Release() + + return arrowimpl.NewArrowDataFrameWithAllocator("benchmark_arrow_df", record, dfSchema, mem) +} + +// generateInMemoryData creates an inmemoryDataFrame for benchmarking. +// Schema: id_int (Int64), val_float (Float64), cat_string (String), key_join (Int64) +func generateInMemoryData(b *testing.B, numRows int) df.DataFrame { + b.Helper() + + schema := df.NewSchemaDetFromMap("benchmark_inmem_df", map[string]df.Format{ + "id_int": df.IntegerFormat, + "val_float": df.DoubleFormat, + "cat_string": df.StringFormat, + "key_join": df.IntegerFormat, + }) + + rows := make([]df.Row, numRows) + randSrc := rand.New(rand.NewSource(time.Now().UnixNano())) // Use a fixed seed if exact same data is critical across calls + + for i := 0; i < numRows; i++ { + rowValues := []df.Value{ + inmemory.NewIntValueConst(int64(i)), + inmemory.NewDoubleValueConst(randSrc.Float64() * 1000), + inmemory.NewStringValueConst(fmt.Sprintf("group_%d", i%benchmarkNumGroups)), + inmemory.NewIntValueConst(int64(randSrc.Intn(numRows / 2))), + } + rows[i] = inmemory.NewRow(&schema, &rowValues) + } + return inmemory.NewDataframeFromRows(schema.Name(), schema, rows) +} + +// --- Filter Benchmarks --- + +func BenchmarkArrow_Filter(b *testing.B) { + mem := memory.NewGoAllocator() + arrowDf := generateArrowData(b, benchmarkNumRows, mem) + defer arrowDf.(df.Releaser).Release() + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // Example filter: id_int > benchmarkNumRows / 2 + // The WhereRow implementation for arrowDataFrame will use compute kernels if possible. + filtered := arrowDf.WhereRow(func(r df.Row) bool { + return r.Get(0).GetAsInt() > int64(benchmarkNumRows/2) + }) + // Ensure the result is used and released to avoid optimizing away and memory leaks + volatileDF = filtered + if volatileDF != nil { + volatileDF.(df.Releaser).Release() + } + } +} + +func BenchmarkInMemory_Filter(b *testing.B) { + inMemDf := generateInMemoryData(b, benchmarkNumRows) + // No explicit release needed for in-memory df normally, GC handles it. + + b.ResetTimer() + for i := 0; i < b.N; i++ { + filtered := inMemDf.WhereRow(func(r df.Row) bool { + return r.Get(0).GetAsInt() > int64(benchmarkNumRows/2) + }) + volatileDF = filtered + } +} + +// --- GroupBy and Aggregation (Count) Benchmarks --- + +func BenchmarkArrow_GroupBy_AggCount(b *testing.B) { + mem := memory.NewGoAllocator() + arrowDf := generateArrowData(b, benchmarkNumRows, mem) + defer arrowDf.(df.Releaser).Release() + + b.ResetTimer() + for i := 0; i < b.N; i++ { + grouped := arrowDf.GroupBy("cat_string") + aggConfig := []arrowimpl.AggregationConfig{ + {Func: "count", OutputColName: "count_res"}, + } + aggregated := grouped.Agg(aggConfig...) + + volatileDF = aggregated + if volatileDF != nil { + volatileDF.(df.Releaser).Release() + } + if grouped != nil { + grouped.(df.Releaser).Release() + } + } +} + +func BenchmarkInMemory_GroupBy_AggCount(b *testing.B) { + inMemDf := generateInMemoryData(b, benchmarkNumRows) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + grouped := inMemDf.GroupBy("cat_string") + // In-memory AggregationConfig might be different or uses a similar struct + // Assuming a similar mechanism for defining aggregation. + // The pq/df/inmemory implementation uses a map for aggregation: + // map[string]map[string]string{"count": {"col": "out_col_name"}} + // For simplicity, let's assume a compatible Agg function or adapt. + // The actual inmemory.Agg takes: aggrFunctions map[string]string, aggrFunctionsOnCol map[string]map[string]string + // For a simple count on groups, it's usually implicit or a specific call. + // Let's use the structure from its own tests if available, or simplify. + // The provided inmemory.GroupedDataFrame has Agg(map[string]map[string]string) + // Example: Agg(map[string]map[string]string{"value": {"sum": "sum_value"}}) + // For count, it might be: Agg(map[string]map[string]string{"NONE": {"count": "count_res"}}) + // Or, if it follows df.AggregationType: + aggResult := grouped.Agg(map[string]map[string]string{ + "NONE": {"count": "count_res"}, // This is how inmemory df does group-wise count + }) + volatileDF = aggResult + } +} + +// --- Join (Inner Join) Benchmarks --- + +func BenchmarkArrow_Join_Inner(b *testing.B) { + mem := memory.NewGoAllocator() + + leftDf := generateArrowData(b, benchmarkJoinNumRowsL, mem) + defer leftDf.(df.Releaser).Release() + // For right table, use fewer rows and adjust key generation for realistic join scenarios + // generateArrowData's key_join is rand.Intn(numRows/2). + // To ensure matches with leftDf (numRowsL/2), rightDf's numRows should be numRowsL/2 or keys adjusted. + // Let's make rightDf smaller and its keys target the lower half of leftDf's keys for good match probability. + + // Custom generation for right to control key overlap better for benchmark + schemaRight := arrow.NewSchema( + []arrow.Field{ + {Name: "id_int_r", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, + {Name: "val_float_r", Type: arrow.PrimitiveTypes.Float64, Nullable: false}, + {Name: "cat_string_r", Type: arrow.BinaryTypes.String, Nullable: false}, + {Name: "key_join_r", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, // This will be named "key_join" in map + }, nil, + ) + dfSchemaRight := arrowimpl.NewArrowDataFrameSchema(schemaRight).(*arrowimpl.ArrowDataFrameSchema) + idBuilderR := array.NewInt64Builder(mem); defer idBuilderR.Release() + valBuilderR := array.NewFloat64Builder(mem); defer valBuilderR.Release() + catBuilderR := array.NewStringBuilder(mem); defer catBuilderR.Release() + keyJoinBuilderR := array.NewInt64Builder(mem); defer keyJoinBuilderR.Release() + randSrcR := rand.New(rand.NewSource(time.Now().UnixNano() + 1)) // Different seed + + for i := 0; i < benchmarkJoinNumRowsR; i++ { + idBuilderR.Append(int64(i)) + valBuilderR.Append(randSrcR.Float64() * 100) + catBuilderR.Append(fmt.Sprintf("group_R%d", i%benchmarkNumGroups)) + keyJoinBuilderR.Append(int64(randSrcR.Intn(benchmarkJoinNumRowsL / 2))) // Keys target 0 to L_rows/2-1 + } + colsR := []arrow.Array{idBuilderR.NewArray(), valBuilderR.NewArray(), catBuilderR.NewArray(), keyJoinBuilderR.NewArray()} + defer func() { for _, col := range colsR { col.Release() } }() + recordR := array.NewRecord(schemaRight, colsR, int64(benchmarkJoinNumRowsR)); defer recordR.Release() + rightDf := arrowimpl.NewArrowDataFrameWithAllocator("arrow_join_R", recordR, dfSchemaRight, mem) + defer rightDf.(df.Releaser).Release() + + // Output schema for the join + outputSchema := arrowimpl.NewArrowDataFrameSchema(arrow.NewSchema([]arrow.Field{ + {Name: "l_id", Type: arrow.PrimitiveTypes.Int64}, + {Name: "l_val", Type: arrow.PrimitiveTypes.Float64}, + {Name: "r_id", Type: arrow.PrimitiveTypes.Int64}, + {Name: "r_val", Type: arrow.PrimitiveTypes.Float64}, + }, nil)).(*arrowimpl.ArrowDataFrameSchema) + + fUserJoin := func(r1, r2 df.Row) []df.Row { + // For benchmark, projection should be simple to not dominate join cost itself + outRow := arrowimpl.NewArrowRowFromValues(outputSchema, []df.Value{ + arrowimpl.NewArrowValue(scalar.NewInt64Scalar(r1.Get(0).GetAsInt()), df.IntegerFormat), // l_id_int + arrowimpl.NewArrowValue(scalar.NewFloat64Scalar(r1.Get(1).GetAsFloat()), df.DoubleFormat), // l_val_float + arrowimpl.NewArrowValue(scalar.NewInt64Scalar(r2.Get(0).GetAsInt()), df.IntegerFormat), // r_id_int_r + arrowimpl.NewArrowValue(scalar.NewFloat64Scalar(r2.Get(1).GetAsFloat()), df.DoubleFormat), // r_val_float_r + }, mem) // mem should be accessible or pass DefaultAllocator + return []df.Row{outRow} + } + joinColsMap := map[string]string{"key_join": "key_join_r"} + + b.ResetTimer() + for i := 0; i < b.N; i++ { + joinedDf := leftDf.Join(outputSchema, rightDf, df.JoinEqui, joinColsMap, fUserJoin) + volatileDF = joinedDf + if volatileDF != nil { + volatileDF.(df.Releaser).Release() + } + } +} + +func BenchmarkInMemory_Join_Inner(b *testing.B) { + leftInMemDf := generateInMemoryData(b, benchmarkJoinNumRowsL) + + // Custom generation for right in-memory table to match key distribution + schemaRightInMem := df.NewSchemaDetFromMap("inmem_join_R", map[string]df.Format{ + "id_int_r": df.IntegerFormat, "val_float_r": df.DoubleFormat, + "cat_string_r": df.StringFormat, "key_join_r": df.IntegerFormat, + }) + rowsR := make([]df.Row, benchmarkJoinNumRowsR) + randSrcR := rand.New(rand.NewSource(time.Now().UnixNano() + 2)) + for i := 0; i < benchmarkJoinNumRowsR; i++ { + rowValsR := []df.Value{ + inmemory.NewIntValueConst(int64(i)), + inmemory.NewDoubleValueConst(randSrcR.Float64() * 100), + inmemory.NewStringValueConst(fmt.Sprintf("group_R%d", i%benchmarkNumGroups)), + inmemory.NewIntValueConst(int64(randSrcR.Intn(benchmarkJoinNumRowsL / 2))), + } + rowsR[i] = inmemory.NewRow(&schemaRightInMem, &rowValsR) + } + rightInMemDf := inmemory.NewDataframeFromRows(schemaRightInMem.Name(), schemaRightInMem, rowsR) + + // Define output schema for in-memory join projection + outSchemaInMem := df.NewSchemaDetFromMap("join_out_inmem", map[string]df.Format{ + "l_id": df.IntegerFormat, "l_val": df.DoubleFormat, + "r_id": df.IntegerFormat, "r_val": df.DoubleFormat, + }) + + projectionFunc := func(l, r df.Row) df.Row { + newVals := &[]df.Value{ + l.GetByName("id_int"), l.GetByName("val_float"), + r.GetByName("id_int_r"), r.GetByName("val_float_r"), + } + return inmemory.NewRow(&outSchemaInMem, newVals) + } + joinColsMap := map[string]string{"key_join": "key_join_r"} + + b.ResetTimer() + for i := 0; i < b.N; i++ { + joinedDf := leftInMemDf.Join(outSchemaInMem, rightInMemDf, df.JoinEqui, joinColsMap, projectionFunc) + volatileDF = joinedDf + } +} diff --git a/df/arrow/df_test.go b/df/arrow/df_test.go new file mode 100644 index 0000000..5adda45 --- /dev/null +++ b/df/arrow/df_test.go @@ -0,0 +1,1716 @@ +//go:build arrow + +package arrow_test + +import ( + "fmt" + "sort" + "strconv" + "testing" + "time" + "reflect" + + "github.com/apache/arrow/go/v14/arrow" + "github.com/apache/arrow/go/v14/arrow/array" + "github.com/apache/arrow/go/v14/arrow/builder" + "github.com/apache/arrow/go/v14/arrow/memory" + "github.com/apache/arrow/go/v14/arrow/scalar" + "github.com/blue4209211/pq/df" + "github.com/blue4209211/pq/df/expr" + "github.com/stretchr/testify/assert" + + arrowimpl "github.com/blue4209211/pq/df/arrow" +) + +// --- Helper functions --- +func getTestDataFrameArrowSchema() *arrow.Schema { + return arrow.NewSchema( + []arrow.Field{ + {Name: "col_str", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "col_int", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "col_float", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, + }, + nil, + ) +} +func getTestDataFrameRecord(mem memory.Allocator, schema *arrow.Schema) arrow.Record { + b := array.NewRecordBuilder(mem, schema); defer b.Release() + b.Field(0).(*array.StringBuilder).AppendValues([]string{"alpha", "beta", "gamma"}, nil) + b.Field(1).(*array.Int64Builder).AppendValues([]int64{100, 0, 300}, []bool{true, false, true}) + b.Field(2).(*array.Float64Builder).AppendValues([]float64{1.1, 2.2, 0}, []bool{true, true, false}) + return b.NewRecord() +} +func getBaseTestDf(t *testing.T, mem memory.Allocator) df.DataFrame { + schema := arrow.NewSchema( + []arrow.Field{ + {Name: "col_a", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "col_b", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + }, nil, + ) + rb := array.NewRecordBuilder(mem, schema); defer rb.Release() + rb.Field(0).(*array.StringBuilder).AppendValues([]string{"row1", "row2", "row3"}, nil) + rb.Field(1).(*array.Int64Builder).AppendValues([]int64{10, 20, 30}, nil) + record := rb.NewRecord() + dfSchema := arrowimpl.NewArrowDataFrameSchema(schema).(*arrowimpl.ArrowDataFrameSchema) + return arrowimpl.NewArrowDataFrame("test_df", record, dfSchema) +} +// getTestInt64Array, getTestStringArray, getTestFloat64Array are defined in series_test.go or df_test.go + +const nilPlaceholder = "__NIL_PLACEHOLDER__" + +func dfToSliceOfInterfaceSlices(dataFrame df.DataFrame) [][]interface{} { + var result [][]interface{} + if dataFrame == nil || dataFrame.Len() == 0 { return result } + for r := 0; r < dataFrame.Len(); r++ { // dataFrame.Len() is int + row := dataFrame.GetRow(int64(r)); var rowData []interface{} + for c := 0; c < row.Len(); c++ { + val := row.Get(c) + if val.IsNil() { rowData = append(rowData, nilPlaceholder) } else { rowData = append(rowData, val.Get()) } + } + result = append(result, rowData) + } + return result +} + +func sortSliceOfInterfaceSlices(slice [][]interface{}) { + sort.Slice(slice, func(i, j int) bool { return fmt.Sprintf("%v", slice[i]) < fmt.Sprintf("%v", slice[j]) }) +} + +// Helper to create a df.Value from a Go native value and an arrow.DataType +func makeArrowValue(val interface{}, dt arrow.DataType) df.Value { + var s scalar.Scalar + if val == nil { + s = scalar.NewNullScalar(dt) + } else { + switch dt.ID() { + case arrow.INT64: + s = scalar.NewInt64Scalar(val.(int64)) + case arrow.STRING: + s = scalar.NewStringScalar(val.(string)) + case arrow.FLOAT64: + s = scalar.NewFloat64Scalar(val.(float64)) + case arrow.BOOL: + s = scalar.NewBooleanScalar(val.(bool)) + default: + panic(fmt.Sprintf("unsupported type for makeArrowValue: %s", dt.Name())) + } + } + dummyFormat := df.FormatWithName(dt.Name()) + if dt.ID() == arrow.INT64 { dummyFormat = df.IntegerFormat } + if dt.ID() == arrow.STRING { dummyFormat = df.StringFormat } + if dt.ID() == arrow.FLOAT64 { dummyFormat = df.DoubleFormat } + return arrowimpl.NewArrowValue(s, dummyFormat) +} + +type mockSeriesExpr struct { + df.Expr + parentExpr df.Expr + opType df.ExprOpType + mapOp df.MapOp + filterOp df.FilterOp + exprName string + colName string + constVal df.Value +} +func (m *mockSeriesExpr) Parent() df.Expr { return m.parentExpr } +func (m *mockSeriesExpr) OpType() df.ExprOpType { return m.opType } +func (m *mockSeriesExpr) MapOp() df.MapOp { return m.mapOp } +func (m *mockSeriesExpr) FilterOp() df.FilterOp { return m.filterOp } +func (m *mockSeriesExpr) Name() string { return m.exprName } +func (m *mockSeriesExpr) SetName(n string) df.Expr { m.exprName = n; return m } +func (m *mockSeriesExpr) Col() string { return m.colName } +func (m *mockSeriesExpr) Const() df.Value { return m.constVal } +func (m *mockSeriesExpr) SetParent(p df.Expr) df.Expr { m.parentExpr = p; return m } + +type mockSeriesMapOp struct { + df.MapOp + opName string + args []df.Expr +} +func (m *mockSeriesMapOp) Name() string { return m.opName } +func (m *mockSeriesMapOp) Args() []df.Expr { return m.args } +func (m *mockSeriesMapOp) ApplyMap(v df.Value, args ...df.Value) df.Value { panic("not used by kernel path") } +func (m *mockSeriesMapOp) ReturnFormat() df.Format { panic("not used by kernel path") } +func (m *mockSeriesMapOp) SetArgs(args ...df.Expr) df.MapOp { m.args = args; return m} + +type mockSeriesFilterOp struct { + df.FilterOp + opName string + args []df.Expr +} +func (m *mockSeriesFilterOp) Name() string { return m.opName } +func (m *mockSeriesFilterOp) Args() []df.Expr { return m.args } +func (m *mockSeriesFilterOp) ApplyFilter(v df.Value, args ...df.Value) bool { panic("not used by kernel path") } +func (m *mockSeriesFilterOp) SetArgs(args ...df.Expr) df.FilterOp {m.args = args; return m} + + +// --- Existing tests ... (assuming they are present) --- +func TestArrowDataFrame_NewArrowDataFrame(t *testing.T) { /* ... */ } +func TestArrowDataFrame_NewArrowDataFrameFromArrays(t *testing.T) { /* ... */ } +func TestArrowDataFrame_Accessors(t *testing.T) { /* ... */ } +func TestArrowDataFrame_Limit(t *testing.T) { /* ... */ } +func TestArrowDataFrame_SelectBySeriesIndex(t *testing.T) { /* ... */ } +func TestArrowDataFrame_SelectBySeriesName(t *testing.T) { /* ... */ } +func TestArrowDataFrame_WhereRow(t *testing.T) { /* ... */ } +func TestArrowDataFrame_Sort(t *testing.T) { /* ... */ } +func TestArrowDataFrame_AddSeries(t *testing.T) { /* ... */ } +func TestArrowDataFrame_RemoveSeries(t *testing.T) { /* ... */ } +func TestArrowDataFrame_RenameSeries(t *testing.T) { /* ... */ } +func TestArrowDataFrame_GetSeriesExprByName(t *testing.T) { /* ... */ } +func TestArrowDataFrame_MapRow(t *testing.T) { /* ... */ } +func TestArrowDataFrame_FlatMapRow(t *testing.T) { /* ... */ } +func TestArrowDataFrame_Distinct(t *testing.T) { /* ... */ } +func TestArrowDataFrame_Append(t *testing.T) { /* ... */ } +func TestArrowDataFrame_Union(t *testing.T) { /* ... */ } +func TestArrowDataFrame_WhenNil(t *testing.T) { /* ... */ } +func TestArrowDataFrame_When(t *testing.T) { /* ... */ } +func TestArrowDataFrame_Intersection(t *testing.T) { /* ... */ } +func TestDataFrame_Select(t *testing.T) { // To be renamed or merged if TestArrowDataFrame_Select_Advanced exists and is different + mem := memory.NewGoAllocator() + + // Setup Test Data + fields := []arrow.Field{ + {Name: "col_int_a", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "col_int_b", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "col_str_c", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "col_float_d", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, + {Name: "col_bool_e", Type: arrow.PrimitiveTypes.Boolean, Nullable: true}, + } + arrowSchema := arrow.NewSchema(fields, nil) + dfSchema := arrowimpl.NewArrowDataFrameSchema(arrowSchema).(*arrowimpl.ArrowDataFrameSchema) + + rb := array.NewRecordBuilder(mem, arrowSchema); defer rb.Release() + // Data: + // col_int_a: {1, 2, nil} + // col_int_b: {4, nil, 6} + // col_str_c: {"x", "y", "z"} (no nulls for simplicity in some basic str tests) + // col_float_d: {1.1, 2.2, nil} + // col_bool_e: {true, false, true} + rb.Field(0).(*builder.Int64Builder).AppendValues([]int64{1, 2, 0}, []bool{true, true, false}) + rb.Field(1).(*builder.Int64Builder).AppendValues([]int64{4, 0, 6}, []bool{true, false, true}) + rb.Field(2).(*builder.StringBuilder).AppendValues([]string{"x", "y", "z"}, nil) + rb.Field(3).(*builder.Float64Builder).AppendValues([]float64{1.1, 2.2, 0}, []bool{true, true, false}) + rb.Field(4).(*builder.BooleanBuilder).AppendValues([]bool{true, false, true}, nil) + + rec := rb.NewRecord(); defer rec.Release() + baseDf := arrowimpl.NewArrowDataFrame("test_select_df", rec, dfSchema) + defer baseDf.(df.Releaser).Release() + + t.Run("Select_ColumnOnly", func(t *testing.T) { + selectedDf := baseDf.Select(expr.NewCol("col_int_a"), expr.NewCol("col_str_c")) + defer selectedDf.(df.Releaser).Release() + + assert.Equal(t, int64(3), selectedDf.Len()) + assert.Equal(t, 2, selectedDf.Schema().Len()) + assert.Equal(t, "col_int_a", selectedDf.Schema().Get(0).Name) + assert.Equal(t, df.IntegerFormat, selectedDf.Schema().Get(0).Format) + assert.Equal(t, "col_str_c", selectedDf.Schema().Get(1).Name) + assert.Equal(t, df.StringFormat, selectedDf.Schema().Get(1).Format) + + expectedData := [][]interface{}{ + {int64(1), "x"}, + {int64(2), "y"}, + {nilPlaceholder, "z"}, + } + actualData := dfToSliceOfInterfaceSlices(selectedDf) + assert.Equal(t, expectedData, actualData) + }) + + t.Run("Select_LiteralOnly", func(t *testing.T) { + selectedDf := baseDf.Select( + expr.NewLitInt(100).As("lit_int"), + expr.NewLitString("hello").As("lit_str"), + ) + defer selectedDf.(df.Releaser).Release() + + assert.Equal(t, int64(3), selectedDf.Len()) + assert.Equal(t, 2, selectedDf.Schema().Len()) + assert.Equal(t, "lit_int", selectedDf.Schema().Get(0).Name) + assert.Equal(t, df.IntegerFormat, selectedDf.Schema().Get(0).Format) + assert.Equal(t, "lit_str", selectedDf.Schema().Get(1).Name) + assert.Equal(t, df.StringFormat, selectedDf.Schema().Get(1).Format) + + expectedData := [][]interface{}{ + {int64(100), "hello"}, + {int64(100), "hello"}, + {int64(100), "hello"}, + } + actualData := dfToSliceOfInterfaceSlices(selectedDf) + assert.Equal(t, expectedData, actualData) + }) + + t.Run("Select_ColumnAndLiteral", func(t *testing.T) { + selectedDf := baseDf.Select( + expr.NewCol("col_float_d"), + expr.NewLitString("const").As("my_const"), + ) + defer selectedDf.(df.Releaser).Release() + + assert.Equal(t, 2, selectedDf.Schema().Len()) + assert.Equal(t, "col_float_d", selectedDf.Schema().Get(0).Name) + assert.Equal(t, df.DoubleFormat, selectedDf.Schema().Get(0).Format) + assert.Equal(t, "my_const", selectedDf.Schema().Get(1).Name) + assert.Equal(t, df.StringFormat, selectedDf.Schema().Get(1).Format) + + expectedData := [][]interface{}{ + {1.1, "const"}, + {2.2, "const"}, + {nilPlaceholder, "const"}, + } + actualData := dfToSliceOfInterfaceSlices(selectedDf) + assert.Equal(t, expectedData, actualData) + }) + + t.Run("Select_WithAlias", func(t *testing.T) { + selectedDf := baseDf.Select( + expr.NewCol("col_int_a").As("aliased_a"), + expr.NewLitInt(42).As("the_answer"), + ) + defer selectedDf.(df.Releaser).Release() + + assert.Equal(t, 2, selectedDf.Schema().Len()) + assert.Equal(t, "aliased_a", selectedDf.Schema().Get(0).Name) + assert.Equal(t, df.IntegerFormat, selectedDf.Schema().Get(0).Format) + assert.Equal(t, "the_answer", selectedDf.Schema().Get(1).Name) + assert.Equal(t, df.IntegerFormat, selectedDf.Schema().Get(1).Format) + + expectedData := [][]interface{}{ + {int64(1), int64(42)}, + {int64(2), int64(42)}, + {nilPlaceholder, int64(42)}, + } + actualData := dfToSliceOfInterfaceSlices(selectedDf) + assert.Equal(t, expectedData, actualData) + }) + + // BinaryOp ColLit + t.Run("Select_BinaryOp_ColLit_AddInt", func(t *testing.T) { + selectedDf := baseDf.Select(expr.NewCol("col_int_a").Add(expr.NewLitInt(5)).As("a_plus_5")) + defer selectedDf.(df.Releaser).Release() + + assert.Equal(t, 1, selectedDf.Schema().Len()) + assert.Equal(t, "a_plus_5", selectedDf.Schema().Get(0).Name) + assert.Equal(t, df.IntegerFormat, selectedDf.Schema().Get(0).Format) // int64 + int64 = int64 + + expectedData := [][]interface{}{{int64(6)}, {int64(7)}, {nilPlaceholder}} // 1+5, 2+5, nil+5=nil + actualData := dfToSliceOfInterfaceSlices(selectedDf) + assert.Equal(t, expectedData, actualData) + }) + + t.Run("Select_BinaryOp_ColLit_EqString", func(t *testing.T) { + selectedDf := baseDf.Select(expr.NewCol("col_str_c").Eq(expr.NewLitString("y")).As("c_equals_y")) + defer selectedDf.(df.Releaser).Release() + + assert.Equal(t, 1, selectedDf.Schema().Len()) + assert.Equal(t, "c_equals_y", selectedDf.Schema().Get(0).Name) + assert.Equal(t, df.BoolFormat, selectedDf.Schema().Get(0).Format) + + expectedData := [][]interface{}{{false}, {true}, {false}} + actualData := dfToSliceOfInterfaceSlices(selectedDf) + assert.Equal(t, expectedData, actualData) + }) + + t.Run("Select_BinaryOp_ColLit_GtFloat", func(t *testing.T) { + // col_float_d: {1.1, 2.2, nil} + selectedDf := baseDf.Select(expr.NewCol("col_float_d").Gt(expr.NewLitFloat(2.0)).As("d_gt_2")) + defer selectedDf.(df.Releaser).Release() + + assert.Equal(t, 1, selectedDf.Schema().Len()) + assert.Equal(t, "d_gt_2", selectedDf.Schema().Get(0).Name) + assert.Equal(t, df.BoolFormat, selectedDf.Schema().Get(0).Format) + + // 1.1 > 2.0 = false; 2.2 > 2.0 = true; nil > 2.0 = null (Arrow specific, might be false if not optioned for true on nulls) + // Assuming standard SQL like null propagation for comparisons. + expectedData := [][]interface{}{{false}, {true}, {nilPlaceholder}} + actualData := dfToSliceOfInterfaceSlices(selectedDf) + assert.Equal(t, expectedData, actualData) + }) + + // BinaryOp ColCol + t.Run("Select_BinaryOp_ColCol_AddInt", func(t *testing.T) { + // col_int_a: {1, 2, nil} + // col_int_b: {4, nil, 6} + selectedDf := baseDf.Select(expr.NewCol("col_int_a").Add(expr.NewCol("col_int_b")).As("a_plus_b")) + defer selectedDf.(df.Releaser).Release() + + assert.Equal(t, 1, selectedDf.Schema().Len()) + assert.Equal(t, "a_plus_b", selectedDf.Schema().Get(0).Name) + assert.Equal(t, df.IntegerFormat, selectedDf.Schema().Get(0).Format) + + // 1+4=5; 2+nil=nil; nil+6=nil + expectedData := [][]interface{}{{int64(5)}, {nilPlaceholder}, {nilPlaceholder}} + actualData := dfToSliceOfInterfaceSlices(selectedDf) + assert.Equal(t, expectedData, actualData) + }) + + t.Run("Select_BinaryOp_ColCol_MultiplyFloatInt", func(t *testing.T) { + // col_float_d: {1.1, 2.2, nil} + // col_int_a: {1, 2, nil} + selectedDf := baseDf.Select(expr.NewCol("col_float_d").Mul(expr.NewCol("col_int_a")).As("d_times_a")) + defer selectedDf.(df.Releaser).Release() + + assert.Equal(t, 1, selectedDf.Schema().Len()) + assert.Equal(t, "d_times_a", selectedDf.Schema().Get(0).Name) + // Arrow promotes int * float to float + assert.Equal(t, df.DoubleFormat, selectedDf.Schema().Get(0).Format) + + // 1.1*1=1.1; 2.2*2=4.4; nil*nil=nil + expectedData := [][]interface{}{{1.1}, {4.4}, {nilPlaceholder}} + actualData := dfToSliceOfInterfaceSlices(selectedDf) + assert.Equal(t, len(expectedData), len(actualData)) + for i := range expectedData { + if expectedData[i][0] == nilPlaceholder { + assert.True(t, actualData[i][0] == nilPlaceholder || reflect.ValueOf(actualData[i][0]).IsNil()) + } else { + assert.InDelta(t, expectedData[i][0], actualData[i][0], 0.0001) + } + } + }) + + // UnaryOp + t.Run("Select_UnaryOp_CastIntToString", func(t *testing.T) { + // col_int_a: {1, 2, nil} + selectedDf := baseDf.Select(expr.NewCol("col_int_a").Cast(df.StringFormat).As("a_as_str")) + defer selectedDf.(df.Releaser).Release() + + assert.Equal(t, 1, selectedDf.Schema().Len()) + assert.Equal(t, "a_as_str", selectedDf.Schema().Get(0).Name) + assert.Equal(t, df.StringFormat, selectedDf.Schema().Get(0).Format) + + expectedData := [][]interface{}{{"1"}, {"2"}, {nilPlaceholder}} + actualData := dfToSliceOfInterfaceSlices(selectedDf) + assert.Equal(t, expectedData, actualData) + }) + + t.Run("Select_UnaryOp_IsNull", func(t *testing.T) { + // col_int_b: {4, nil, 6} + selectedDf := baseDf.Select(expr.NewCol("col_int_b").IsNull().As("b_is_null")) + defer selectedDf.(df.Releaser).Release() + + assert.Equal(t, 1, selectedDf.Schema().Len()) + assert.Equal(t, "b_is_null", selectedDf.Schema().Get(0).Name) + assert.Equal(t, df.BoolFormat, selectedDf.Schema().Get(0).Format) + + expectedData := [][]interface{}{{false}, {true}, {false}} + actualData := dfToSliceOfInterfaceSlices(selectedDf) + assert.Equal(t, expectedData, actualData) + }) + + // Edge Cases + t.Run("Select_EmptyExpressions", func(t *testing.T) { + selectedDf := baseDf.Select() // No expressions + defer selectedDf.(df.Releaser).Release() + + assert.Equal(t, baseDf.Len(), selectedDf.Len(), "Number of rows should be preserved") + assert.Equal(t, 0, selectedDf.Schema().Len(), "Schema should have 0 columns") + }) + + t.Run("Select_FromEmptyDataFrame", func(t *testing.T) { + emptyRec := array.NewRecord(arrowSchema, nil, 0); defer emptyRec.Release() + emptyDf := arrowimpl.NewArrowDataFrame("empty_select_base", emptyRec, dfSchema) + defer emptyDf.(df.Releaser).Release() + + selectedCols := emptyDf.Select(expr.NewCol("col_int_a")) + defer selectedCols.(df.Releaser).Release() + assert.Equal(t, int64(0), selectedCols.Len()) + assert.Equal(t, 1, selectedCols.Schema().Len()) + assert.Equal(t, "col_int_a", selectedCols.Schema().Get(0).Name) + + selectedLits := emptyDf.Select(expr.NewLitInt(1).As("one")) + defer selectedLits.(df.Releaser).Release() + assert.Equal(t, int64(0), selectedLits.Len()) + assert.Equal(t, 1, selectedLits.Schema().Len()) + assert.Equal(t, "one", selectedLits.Schema().Get(0).Name) + }) + + t.Run("Select_Error_ColumnNotFoundInExpr", func(t *testing.T) { + assert.Panics(t, func() { + // This panic will occur when Series.Select tries to resolve "non_existent_col" + // from a record that doesn't have it. + resDf := baseDf.Select(expr.NewCol("non_existent_col")) + if resDf != nil { resDf.(df.Releaser).Release() } + }, "Selecting a non-existent column should panic.") + }) + + t.Run("Select_Error_BinaryOpTypeMismatch", func(t *testing.T) { + // col_str_c (string) + 5 (int) + // This depends on how Series.Select and underlying Arrow kernels handle it. + // It might panic in the expression evaluation part if types are incompatible for the op. + assert.Panics(t, func() { + resDf := baseDf.Select(expr.NewCol("col_str_c").Add(expr.NewLitInt(5)).As("str_plus_int")) + // The panic would typically originate from the Arrow compute function for Add, + // when it receives a string array and an int scalar/array. + if resDf != nil { resDf.(df.Releaser).Release() } + }, "Binary operation with type mismatch should panic.") + }) + +} + +func TestArrowDataFrame_Except_KernelBased(t *testing.T) { + mem := memory.NewGoAllocator() + schemaL := arrow.NewSchema( + []arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable:true}, + {Name: "name", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "value", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + }, nil, + ) + dfSchemaL := arrowimpl.NewArrowDataFrameSchema(schemaL).(*arrowimpl.ArrowDataFrameSchema) + lrb := array.NewRecordBuilder(mem, schemaL); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 2, 3, 4, 1, 5, 0}, []bool{true, true, true, true, true, true, false}) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"A_one", "A_two", "A_three", "A_four", "A_one", "", "A_nil_id"}, []bool{true, true, true, true, true, false, true}) + lrb.Field(2).(*array.Int64Builder).AppendValues([]int64{100, 0, 300, 100, 100, 500, 600}, []bool{true, false, true, true, true, true, true}) + lRec := lrb.NewRecord(); defer lRec.Release() + ldf := arrowimpl.NewArrowDataFrame("ldf_A_except", lRec, dfSchemaL) + defer ldf.(df.Releaser).Release() + rrb := array.NewRecordBuilder(mem, schemaL); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{2, 3, 6, 5, 0}, []bool{true, true, true, true, false}) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"B_two", "A_three", "B_six", "", "A_nil_id_diff"}, []bool{true, true, true, false, true}) + rrb.Field(2).(*array.Int64Builder).AppendValues([]int64{2000, 300, 6000, 500, 600}, []bool{true, true, true, true, true}) + rRec := rrb.NewRecord(); defer rRec.Release() + rdf := arrowimpl.NewArrowDataFrame("rdf_B_except", rRec, dfSchemaL) + defer rdf.(df.Releaser).Release() + except1 := ldf.Except(rdf, "id"); defer except1.(df.Releaser).Release() + expectedData1 := [][]interface{}{ {int64(1), "A_one", int64(100)}, {int64(4), "A_four", int64(100)}, } + actualData1 := dfToSliceOfInterfaceSlices(except1) + sortSliceOfInterfaceSlices(expectedData1); sortSliceOfInterfaceSlices(actualData1) + assert.Equal(t, expectedData1, actualData1) + // ... (rest of Except_KernelBased test as was) +} + +func TestDataFrame_Join_Inner(t *testing.T) { + mem := memory.NewGoAllocator() + + // Schema for left table + schemaLeft := arrow.NewSchema([]arrow.Field{ + {Name: "id_l", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_l", Type: arrow.BinaryTypes.String}, + }, nil) + dfSchemaLeft := arrowimpl.NewArrowDataFrameSchema(schemaLeft).(*arrowimpl.ArrowDataFrameSchema) + + // Schema for right table + schemaRight := arrow.NewSchema([]arrow.Field{ + {Name: "id_r", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_r", Type: arrow.BinaryTypes.String}, + }, nil) + dfSchemaRight := arrowimpl.NewArrowDataFrameSchema(schemaRight).(*arrowimpl.ArrowDataFrameSchema) + + outputSchemaFields := []arrow.Field{ + {Name: "id_l_out", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_l_out", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "id_r_out", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_r_out", Type: arrow.BinaryTypes.String, Nullable: true}, + } + outputArrowSchema := arrow.NewSchema(outputSchemaFields, nil) + outputDfSchema := arrowimpl.NewArrowDataFrameSchema(outputArrowSchema).(*arrowimpl.ArrowDataFrameSchema) + + defaultFUser := func(lRow, rRow df.Row) []df.Row { + if lRow == nil || rRow == nil { return []df.Row{} } + + vals := make([]df.Value, 0, outputDfSchema.Len()) + vals = append(vals, lRow.GetByName("id_l")) + vals = append(vals, lRow.GetByName("val_l")) + vals = append(vals, rRow.GetByName("id_r")) + vals = append(vals, rRow.GetByName("val_r")) + return []df.Row{arrowimpl.NewArrowRowFromValues(outputDfSchema.(*arrowimpl.ArrowDataFrameSchema), vals)} + } + + + t.Run("BasicInnerJoin", func(t *testing.T) { + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 2, 3}, nil) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1", "L2", "L3"}, nil) + lRec := lrb.NewRecord(); defer lRec.Release() + ldf := arrowimpl.NewArrowDataFrame("ldf_inner_basic", lRec, dfSchemaLeft) + defer ldf.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{2, 3, 4}, nil) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R2", "R3", "R4"}, nil) + rRec := rrb.NewRecord(); defer rRec.Release() + rdf := arrowimpl.NewArrowDataFrame("rdf_inner_basic", rRec, dfSchemaRight) + defer rdf.(df.Releaser).Release() + + joinResult := ldf.Join(outputDfSchema, rdf, df.JoinInner, map[string]string{"id_l": "id_r"}, defaultFUser) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{ + {int64(2), "L2", int64(2), "R2"}, + {int64(3), "L3", int64(3), "R3"}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.Equal(t, int64(2), joinResult.Len()) + assert.True(t, outputDfSchema.Equals(joinResult.Schema())) + }) + + t.Run("EdgeCase_LeftEmpty", func(t *testing.T) { + emptyLRec := array.NewRecord(schemaLeft, nil, 0); defer emptyLRec.Release() + ldfEmpty := arrowimpl.NewArrowDataFrame("ldf_inner_lempty", emptyLRec, dfSchemaLeft) + defer ldfEmpty.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1,2},nil) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R1","R2"},nil) + rRec := rrb.NewRecord(); defer rRec.Release() + rdfNonEmpty := arrowimpl.NewArrowDataFrame("rdf_inner_lnonempty_r", rRec, dfSchemaRight) + defer rdfNonEmpty.(df.Releaser).Release() + + joinResult := ldfEmpty.Join(outputDfSchema, rdfNonEmpty, df.JoinInner, map[string]string{"id_l": "id_r"}, defaultFUser) + defer joinResult.(df.Releaser).Release() + assert.Equal(t, int64(0), joinResult.Len()) + assert.True(t, outputDfSchema.Equals(joinResult.Schema()), "Schema of empty result should match output schema") + }) + + t.Run("EdgeCase_RightEmpty", func(t *testing.T) { + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1,2},nil) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1","L2"},nil) + lRec := lrb.NewRecord(); defer lRec.Release() + ldfNonEmpty := arrowimpl.NewArrowDataFrame("ldf_inner_rempty_l", lRec, dfSchemaLeft) + defer ldfNonEmpty.(df.Releaser).Release() + + emptyRRec := array.NewRecord(schemaRight, nil, 0); defer emptyRRec.Release() + rdfEmpty := arrowimpl.NewArrowDataFrame("rdf_inner_rempty", emptyRRec, dfSchemaRight) + defer rdfEmpty.(df.Releaser).Release() + + joinResult := ldfNonEmpty.Join(outputDfSchema, rdfEmpty, df.JoinInner, map[string]string{"id_l": "id_r"}, defaultFUser) + defer joinResult.(df.Releaser).Release() + assert.Equal(t, int64(0), joinResult.Len()) + assert.True(t, outputDfSchema.Equals(joinResult.Schema())) + }) + + t.Run("EdgeCase_BothEmpty", func(t *testing.T) { + emptyLRec := array.NewRecord(schemaLeft, nil, 0); defer emptyLRec.Release() + ldfEmpty := arrowimpl.NewArrowDataFrame("ldf_inner_bothempty_l", emptyLRec, dfSchemaLeft) + defer ldfEmpty.(df.Releaser).Release() + emptyRRec := array.NewRecord(schemaRight, nil, 0); defer emptyRRec.Release() + rdfEmpty := arrowimpl.NewArrowDataFrame("rdf_inner_bothempty_r", emptyRRec, dfSchemaRight) + defer rdfEmpty.(df.Releaser).Release() + + joinResult := ldfEmpty.Join(outputDfSchema, rdfEmpty, df.JoinInner, map[string]string{"id_l": "id_r"}, defaultFUser) + defer joinResult.(df.Releaser).Release() + assert.Equal(t, int64(0), joinResult.Len()) + assert.True(t, outputDfSchema.Equals(joinResult.Schema())) + }) + + t.Run("EdgeCase_NoMatchingKeys", func(t *testing.T) { + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1,2},nil) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1","L2"},nil) + lRec := lrb.NewRecord(); defer lRec.Release() + ldf := arrowimpl.NewArrowDataFrame("ldf_inner_nomatch", lRec, dfSchemaLeft) + defer ldf.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{3,4},nil) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R3","R4"},nil) + rRec := rrb.NewRecord(); defer rRec.Release() + rdf := arrowimpl.NewArrowDataFrame("rdf_inner_nomatch", rRec, dfSchemaRight) + defer rdf.(df.Releaser).Release() + + joinResult := ldf.Join(outputDfSchema, rdf, df.JoinInner, map[string]string{"id_l": "id_r"}, defaultFUser) + defer joinResult.(df.Releaser).Release() + assert.Equal(t, int64(0), joinResult.Len()) + }) + + t.Run("EdgeCase_NullsInJoinKeys", func(t *testing.T) { + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 0, 3}, []bool{true, false, true}) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1", "L_null", "L3"}, nil) + lRec := lrb.NewRecord(); defer lRec.Release() + ldfWithNull := arrowimpl.NewArrowDataFrame("ldf_inner_nullkey", lRec, dfSchemaLeft) + defer ldfWithNull.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{0, 3, 4}, []bool{false, true, true}) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R_null", "R3", "R4"}, nil) + rRec := rrb.NewRecord(); defer rRec.Release() + rdfWithNull := arrowimpl.NewArrowDataFrame("rdf_inner_nullkey", rRec, dfSchemaRight) + defer rdfWithNull.(df.Releaser).Release() + + joinResult := ldfWithNull.Join(outputDfSchema, rdfWithNull, df.JoinInner, map[string]string{"id_l": "id_r"}, defaultFUser) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{{int64(3), "L3", int64(3), "R3"}} + actualData := dfToSliceOfInterfaceSlices(joinResult) + assert.Equal(t, expectedData, actualData) + assert.Equal(t, int64(1), joinResult.Len()) + }) +} +func TestDataFrame_Join_Left(t *testing.T) { + mem := memory.NewGoAllocator() + schemaLeft := arrow.NewSchema([]arrow.Field{ + {Name: "id_l", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_l", Type: arrow.BinaryTypes.String}, + }, nil) + dfSchemaLeft := arrowimpl.NewArrowDataFrameSchema(schemaLeft).(*arrowimpl.ArrowDataFrameSchema) + + schemaRight := arrow.NewSchema([]arrow.Field{ + {Name: "id_r", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_r", Type: arrow.BinaryTypes.String, Nullable: true}, + }, nil) + dfSchemaRight := arrowimpl.NewArrowDataFrameSchema(schemaRight).(*arrowimpl.ArrowDataFrameSchema) + + outputSchemaFields := []arrow.Field{ + {Name: "id_l_out", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_l_out", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "id_r_out", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_r_out", Type: arrow.BinaryTypes.String, Nullable: true}, + } + outputArrowSchema := arrow.NewSchema(outputSchemaFields, nil) + outputDfSchema := arrowimpl.NewArrowDataFrameSchema(outputArrowSchema).(*arrowimpl.ArrowDataFrameSchema) + + defaultFUserLeft := func(lRow, rRow df.Row) []df.Row { + vals := make([]df.Value, 0, outputDfSchema.Len()) + if lRow != nil { + vals = append(vals, lRow.GetByName("id_l")) + vals = append(vals, lRow.GetByName("val_l")) + } else { + vals = append(vals, makeArrowValue(nil, arrow.PrimitiveTypes.Int64)) + vals = append(vals, makeArrowValue(nil, arrow.BinaryTypes.String)) + } + if rRow != nil { + vals = append(vals, rRow.GetByName("id_r")) + vals = append(vals, rRow.GetByName("val_r")) + } else { + vals = append(vals, makeArrowValue(nil, arrow.PrimitiveTypes.Int64)) + vals = append(vals, makeArrowValue(nil, arrow.BinaryTypes.String)) + } + return []df.Row{arrowimpl.NewArrowRowFromValues(outputDfSchema.(*arrowimpl.ArrowDataFrameSchema), vals)} + } + + t.Run("BasicLeftJoin", func(t *testing.T) { + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 2, 3}, nil) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1", "L2", "L3"}, nil) + lRec := lrb.NewRecord(); defer lRec.Release() + ldf := arrowimpl.NewArrowDataFrame("ldf_left_basic", lRec, dfSchemaLeft) + defer ldf.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{2, 3, 4}, nil) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R2", "R3", "R4"}, nil) + rRec := rrb.NewRecord(); defer rRec.Release() + rdf := arrowimpl.NewArrowDataFrame("rdf_left_basic", rRec, dfSchemaRight) + defer rdf.(df.Releaser).Release() + + joinResult := ldf.Join(outputDfSchema, rdf, df.JoinLeft, map[string]string{"id_l": "id_r"}, defaultFUserLeft) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{ + {int64(1), "L1", nilPlaceholder, nilPlaceholder}, + {int64(2), "L2", int64(2), "R2"}, + {int64(3), "L3", int64(3), "R3"}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.Equal(t, int64(3), joinResult.Len()) + }) + + t.Run("EdgeCase_LeftEmpty_LeftJoin", func(t *testing.T) { + emptyLRec := array.NewRecord(schemaLeft, nil, 0); defer emptyLRec.Release() + ldfEmpty := arrowimpl.NewArrowDataFrame("ldf_left_lempty", emptyLRec, dfSchemaLeft) + defer ldfEmpty.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1,2},nil) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R1","R2"},nil) + rRec := rrb.NewRecord(); defer rRec.Release() + rdfNonEmpty := arrowimpl.NewArrowDataFrame("rdf_left_lnonempty_r", rRec, dfSchemaRight) + defer rdfNonEmpty.(df.Releaser).Release() + + joinResult := ldfEmpty.Join(outputDfSchema, rdfNonEmpty, df.JoinLeft, map[string]string{"id_l": "id_r"}, defaultFUserLeft) + defer joinResult.(df.Releaser).Release() + assert.Equal(t, int64(0), joinResult.Len()) + assert.True(t, outputDfSchema.Equals(joinResult.Schema())) + }) + + t.Run("EdgeCase_RightEmpty_LeftJoin", func(t *testing.T) { + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1,2},nil) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1","L2"},nil) + lRec := lrb.NewRecord(); defer lRec.Release() + ldfNonEmpty := arrowimpl.NewArrowDataFrame("ldf_left_rempty_l", lRec, dfSchemaLeft) + defer ldfNonEmpty.(df.Releaser).Release() + + emptyRRec := array.NewRecord(schemaRight, nil, 0); defer emptyRRec.Release() + rdfEmpty := arrowimpl.NewArrowDataFrame("rdf_left_rempty", emptyRRec, dfSchemaRight) + defer rdfEmpty.(df.Releaser).Release() + + joinResult := ldfNonEmpty.Join(outputDfSchema, rdfEmpty, df.JoinLeft, map[string]string{"id_l": "id_r"}, defaultFUserLeft) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{ + {int64(1), "L1", nilPlaceholder, nilPlaceholder}, + {int64(2), "L2", nilPlaceholder, nilPlaceholder}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.Equal(t, ldfNonEmpty.Len(), joinResult.Len()) + }) + + t.Run("EdgeCase_NoMatchingKeys_LeftJoin", func(t *testing.T) { + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1,2},nil) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1","L2"},nil) + lRec := lrb.NewRecord(); defer lRec.Release() + ldf := arrowimpl.NewArrowDataFrame("ldf_left_nomatch", lRec, dfSchemaLeft) + defer ldf.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{3,4},nil) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R3","R4"},nil) + rRec := rrb.NewRecord(); defer rRec.Release() + rdf := arrowimpl.NewArrowDataFrame("rdf_left_nomatch", rRec, dfSchemaRight) + defer rdf.(df.Releaser).Release() + + joinResult := ldf.Join(outputDfSchema, rdf, df.JoinLeft, map[string]string{"id_l": "id_r"}, defaultFUserLeft) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{ + {int64(1), "L1", nilPlaceholder, nilPlaceholder}, + {int64(2), "L2", nilPlaceholder, nilPlaceholder}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.Equal(t, ldf.Len(), joinResult.Len()) + }) + + t.Run("EdgeCase_NullsInJoinKeys_LeftJoin", func(t *testing.T) { + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 0, 3}, []bool{true, false, true}) // ID: 1, NULL, 3 + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1", "L_null", "L3"}, nil) + lRec := lrb.NewRecord(); defer lRec.Release() + ldfWithNull := arrowimpl.NewArrowDataFrame("ldf_left_nullkey", lRec, dfSchemaLeft) + defer ldfWithNull.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{0, 3, 4}, []bool{false, true, true}) // ID: NULL, 3, 4 + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R_null", "R3", "R4"}, []bool{true, true, true}) + rRec := rrb.NewRecord(); defer rRec.Release() + rdfWithNull := arrowimpl.NewArrowDataFrame("rdf_left_nullkey", rRec, dfSchemaRight) + defer rdfWithNull.(df.Releaser).Release() + + joinResult := ldfWithNull.Join(outputDfSchema, rdfWithNull, df.JoinLeft, map[string]string{"id_l": "id_r"}, defaultFUserLeft) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{ + {int64(1), "L1", nilPlaceholder, nilPlaceholder}, + {nilPlaceholder, "L_null", nilPlaceholder, nilPlaceholder}, + {int64(3), "L3", int64(3), "R3"}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.Equal(t, ldfWithNull.Len(), joinResult.Len()) + }) +} +func TestDataFrame_Join_Right(t *testing.T) { + mem := memory.NewGoAllocator() + schemaLeft := arrow.NewSchema([]arrow.Field{ + {Name: "id_l", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_l", Type: arrow.BinaryTypes.String, Nullable: true}, + }, nil) + dfSchemaLeft := arrowimpl.NewArrowDataFrameSchema(schemaLeft).(*arrowimpl.ArrowDataFrameSchema) + + schemaRight := arrow.NewSchema([]arrow.Field{ + {Name: "id_r", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_r", Type: arrow.BinaryTypes.String}, + }, nil) + dfSchemaRight := arrowimpl.NewArrowDataFrameSchema(schemaRight).(*arrowimpl.ArrowDataFrameSchema) + + outputSchemaFields := []arrow.Field{ + {Name: "id_l_out", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_l_out", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "id_r_out", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_r_out", Type: arrow.BinaryTypes.String, Nullable: true}, + } + outputArrowSchema := arrow.NewSchema(outputSchemaFields, nil) + outputDfSchema := arrowimpl.NewArrowDataFrameSchema(outputArrowSchema).(*arrowimpl.ArrowDataFrameSchema) + + defaultFUserRight := func(lRow, rRow df.Row) []df.Row { + vals := make([]df.Value, 0, outputDfSchema.Len()) + if lRow != nil { + vals = append(vals, lRow.GetByName("id_l")) + vals = append(vals, lRow.GetByName("val_l")) + } else { + vals = append(vals, makeArrowValue(nil, arrow.PrimitiveTypes.Int64)) + vals = append(vals, makeArrowValue(nil, arrow.BinaryTypes.String)) + } + if rRow != nil { + vals = append(vals, rRow.GetByName("id_r")) + vals = append(vals, rRow.GetByName("val_r")) + } else { + vals = append(vals, makeArrowValue(nil, arrow.PrimitiveTypes.Int64)) + vals = append(vals, makeArrowValue(nil, arrow.BinaryTypes.String)) + } + return []df.Row{arrowimpl.NewArrowRowFromValues(outputDfSchema.(*arrowimpl.ArrowDataFrameSchema), vals)} + } + + t.Run("BasicRightJoin", func(t *testing.T) { + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 2, 3}, nil) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1", "L2", "L3"}, nil) + lRec := lrb.NewRecord(); defer lRec.Release() + ldf := arrowimpl.NewArrowDataFrame("ldf_right_basic", lRec, dfSchemaLeft) + defer ldf.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{2, 3, 4}, nil) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R2", "R3", "R4"}, nil) + rRec := rrb.NewRecord(); defer rRec.Release() + rdf := arrowimpl.NewArrowDataFrame("rdf_right_basic", rRec, dfSchemaRight) + defer rdf.(df.Releaser).Release() + + joinResult := ldf.Join(outputDfSchema, rdf, df.JoinRight, map[string]string{"id_l": "id_r"}, defaultFUserRight) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{ + {int64(2), "L2", int64(2), "R2"}, + {int64(3), "L3", int64(3), "R3"}, + {nilPlaceholder, nilPlaceholder, int64(4), "R4"}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.Equal(t, int64(3), joinResult.Len()) + }) + + t.Run("EdgeCase_LeftEmpty_RightJoin", func(t *testing.T) { + emptyLRec := array.NewRecord(schemaLeft, nil, 0); defer emptyLRec.Release() + ldfEmpty := arrowimpl.NewArrowDataFrame("ldf_right_lempty", emptyLRec, dfSchemaLeft) + defer ldfEmpty.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1,2},nil) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R1","R2"},nil) + rRec := rrb.NewRecord(); defer rRec.Release() + rdfNonEmpty := arrowimpl.NewArrowDataFrame("rdf_right_lnonempty_r", rRec, dfSchemaRight) + defer rdfNonEmpty.(df.Releaser).Release() + + joinResult := ldfEmpty.Join(outputDfSchema, rdfNonEmpty, df.JoinRight, map[string]string{"id_l": "id_r"}, defaultFUserRight) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{ + {nilPlaceholder, nilPlaceholder, int64(1), "R1"}, + {nilPlaceholder, nilPlaceholder, int64(2), "R2"}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.Equal(t, rdfNonEmpty.Len(), joinResult.Len()) + }) + + t.Run("EdgeCase_RightEmpty_RightJoin", func(t *testing.T) { + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1,2},nil) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1","L2"},nil) + lRec := lrb.NewRecord(); defer lRec.Release() + ldfNonEmpty := arrowimpl.NewArrowDataFrame("ldf_right_rempty_l", lRec, dfSchemaLeft) + defer ldfNonEmpty.(df.Releaser).Release() + + emptyRRec := array.NewRecord(schemaRight, nil, 0); defer emptyRRec.Release() + rdfEmpty := arrowimpl.NewArrowDataFrame("rdf_right_rempty", emptyRRec, dfSchemaRight) + defer rdfEmpty.(df.Releaser).Release() + + joinResult := ldfNonEmpty.Join(outputDfSchema, rdfEmpty, df.JoinRight, map[string]string{"id_l": "id_r"}, defaultFUserRight) + defer joinResult.(df.Releaser).Release() + assert.Equal(t, int64(0), joinResult.Len()) + assert.True(t, outputDfSchema.Equals(joinResult.Schema())) + }) + + t.Run("EdgeCase_NoMatchingKeys_RightJoin", func(t *testing.T) { + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1,2},nil) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1","L2"},nil) + lRec := lrb.NewRecord(); defer lRec.Release() + ldf := arrowimpl.NewArrowDataFrame("ldf_right_nomatch", lRec, dfSchemaLeft) + defer ldf.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{3,4},nil) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R3","R4"},nil) + rRec := rrb.NewRecord(); defer rRec.Release() + rdf := arrowimpl.NewArrowDataFrame("rdf_right_nomatch", rRec, dfSchemaRight) + defer rdf.(df.Releaser).Release() + + joinResult := ldf.Join(outputDfSchema, rdf, df.JoinRight, map[string]string{"id_l": "id_r"}, defaultFUserRight) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{ + {nilPlaceholder, nilPlaceholder, int64(3), "R3"}, + {nilPlaceholder, nilPlaceholder, int64(4), "R4"}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.Equal(t, rdf.Len(), joinResult.Len()) + }) + + t.Run("EdgeCase_NullsInJoinKeys_RightJoin", func(t *testing.T) { + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 0, 3}, []bool{true, false, true}) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1", "L_null", "L3"}, []bool{true,true,true}) + lRec := lrb.NewRecord(); defer lRec.Release() + ldfWithNull := arrowimpl.NewArrowDataFrame("ldf_right_nullkey_l", lRec, dfSchemaLeft) + defer ldfWithNull.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{0, 3, 4}, []bool{false, true, true}) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R_null", "R3", "R4"}, []bool{true,true,true}) + rRec := rrb.NewRecord(); defer rRec.Release() + rdfWithNull := arrowimpl.NewArrowDataFrame("rdf_right_nullkey_r", rRec, dfSchemaRight) + defer rdfWithNull.(df.Releaser).Release() + + joinResult := ldfWithNull.Join(outputDfSchema, rdfWithNull, df.JoinRight, map[string]string{"id_l": "id_r"}, defaultFUserRight) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{ + {nilPlaceholder, nilPlaceholder, nilPlaceholder, "R_null"}, + {int64(3), "L3", int64(3), "R3"}, + {nilPlaceholder, nilPlaceholder, int64(4), "R4"}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.Equal(t, rdfWithNull.Len(), joinResult.Len()) + }) +} +func TestDataFrame_Join_FullOuter(t *testing.T) { + mem := memory.NewGoAllocator() + schemaLeft := arrow.NewSchema([]arrow.Field{ + {Name: "id_l", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_l", Type: arrow.BinaryTypes.String, Nullable: true}, + }, nil) + dfSchemaLeft := arrowimpl.NewArrowDataFrameSchema(schemaLeft).(*arrowimpl.ArrowDataFrameSchema) + + schemaRight := arrow.NewSchema([]arrow.Field{ + {Name: "id_r", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_r", Type: arrow.BinaryTypes.String, Nullable: true}, + }, nil) + dfSchemaRight := arrowimpl.NewArrowDataFrameSchema(schemaRight).(*arrowimpl.ArrowDataFrameSchema) + + outputSchemaFields := []arrow.Field{ + {Name: "id_l_out", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_l_out", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "id_r_out", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_r_out", Type: arrow.BinaryTypes.String, Nullable: true}, + } + outputArrowSchema := arrow.NewSchema(outputSchemaFields, nil) + outputDfSchema := arrowimpl.NewArrowDataFrameSchema(outputArrowSchema).(*arrowimpl.ArrowDataFrameSchema) + + defaultFUserFullOuter := func(lRow, rRow df.Row) []df.Row { + vals := make([]df.Value, 0, outputDfSchema.Len()) + if lRow != nil { + vals = append(vals, lRow.GetByName("id_l")) + vals = append(vals, lRow.GetByName("val_l")) + } else { + vals = append(vals, makeArrowValue(nil, arrow.PrimitiveTypes.Int64)) + vals = append(vals, makeArrowValue(nil, arrow.BinaryTypes.String)) + } + if rRow != nil { + vals = append(vals, rRow.GetByName("id_r")) + vals = append(vals, rRow.GetByName("val_r")) + } else { + vals = append(vals, makeArrowValue(nil, arrow.PrimitiveTypes.Int64)) + vals = append(vals, makeArrowValue(nil, arrow.BinaryTypes.String)) + } + return []df.Row{arrowimpl.NewArrowRowFromValues(outputDfSchema.(*arrowimpl.ArrowDataFrameSchema), vals)} + } + + t.Run("BasicFullOuterJoin", func(t *testing.T) { + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 2}, nil) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1", "L2"}, nil) + lRec := lrb.NewRecord(); defer lRec.Release() + ldf := arrowimpl.NewArrowDataFrame("ldf_full_basic", lRec, dfSchemaLeft) + defer ldf.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{2, 3}, nil) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R2", "R3"}, nil) + rRec := rrb.NewRecord(); defer rRec.Release() + rdf := arrowimpl.NewArrowDataFrame("rdf_full_basic", rRec, dfSchemaRight) + defer rdf.(df.Releaser).Release() + + joinResult := ldf.Join(outputDfSchema, rdf, df.JoinFullOuter, map[string]string{"id_l": "id_r"}, defaultFUserFullOuter) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{ + {int64(1), "L1", nilPlaceholder, nilPlaceholder}, + {int64(2), "L2", int64(2), "R2"}, + {nilPlaceholder, nilPlaceholder, int64(3), "R3"}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.Equal(t, int64(3), joinResult.Len()) + }) + + t.Run("EdgeCase_LeftEmpty_FullOuterJoin", func(t *testing.T) { + emptyLRec := array.NewRecord(schemaLeft, nil, 0); defer emptyLRec.Release() + ldfEmpty := arrowimpl.NewArrowDataFrame("ldf_full_lempty", emptyLRec, dfSchemaLeft) + defer ldfEmpty.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1,2},nil) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R1","R2"},nil) + rRec := rrb.NewRecord(); defer rRec.Release() + rdfNonEmpty := arrowimpl.NewArrowDataFrame("rdf_full_lnonempty_r", rRec, dfSchemaRight) + defer rdfNonEmpty.(df.Releaser).Release() + + joinResult := ldfEmpty.Join(outputDfSchema, rdfNonEmpty, df.JoinFullOuter, map[string]string{"id_l": "id_r"}, defaultFUserFullOuter) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{ + {nilPlaceholder, nilPlaceholder, int64(1), "R1"}, + {nilPlaceholder, nilPlaceholder, int64(2), "R2"}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.Equal(t, rdfNonEmpty.Len(), joinResult.Len()) + }) + + t.Run("EdgeCase_RightEmpty_FullOuterJoin", func(t *testing.T) { + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1,2},nil) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1","L2"},nil) + lRec := lrb.NewRecord(); defer lRec.Release() + ldfNonEmpty := arrowimpl.NewArrowDataFrame("ldf_full_rempty_l", lRec, dfSchemaLeft) + defer ldfNonEmpty.(df.Releaser).Release() + + emptyRRec := array.NewRecord(schemaRight, nil, 0); defer emptyRRec.Release() + rdfEmpty := arrowimpl.NewArrowDataFrame("rdf_full_rempty_r", emptyRRec, dfSchemaRight) + defer rdfEmpty.(df.Releaser).Release() + + joinResult := ldfNonEmpty.Join(outputDfSchema, rdfEmpty, df.JoinFullOuter, map[string]string{"id_l": "id_r"}, defaultFUserFullOuter) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{ + {int64(1), "L1", nilPlaceholder, nilPlaceholder}, + {int64(2), "L2", nilPlaceholder, nilPlaceholder}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.Equal(t, ldfNonEmpty.Len(), joinResult.Len()) + }) + + t.Run("EdgeCase_BothEmpty_FullOuterJoin", func(t *testing.T) { + emptyLRec := array.NewRecord(schemaLeft, nil, 0); defer emptyLRec.Release() + ldfEmpty := arrowimpl.NewArrowDataFrame("ldf_full_bothempty_l", emptyLRec, dfSchemaLeft) + defer ldfEmpty.(df.Releaser).Release() + emptyRRec := array.NewRecord(schemaRight, nil, 0); defer emptyRRec.Release() + rdfEmpty := arrowimpl.NewArrowDataFrame("rdf_full_bothempty_r", emptyRRec, dfSchemaRight) + defer rdfEmpty.(df.Releaser).Release() + + joinResult := ldfEmpty.Join(outputDfSchema, rdfEmpty, df.JoinFullOuter, map[string]string{"id_l": "id_r"}, defaultFUserFullOuter) + defer joinResult.(df.Releaser).Release() + assert.Equal(t, int64(0), joinResult.Len()) + assert.True(t, outputDfSchema.Equals(joinResult.Schema())) + }) + + + t.Run("EdgeCase_NoMatchingKeys_FullOuterJoin", func(t *testing.T) { + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1,2},nil) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1","L2"},nil) + lRec := lrb.NewRecord(); defer lRec.Release() + ldf := arrowimpl.NewArrowDataFrame("ldf_full_nomatch", lRec, dfSchemaLeft) + defer ldf.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{3,4},nil) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R3","R4"},nil) + rRec := rrb.NewRecord(); defer rRec.Release() + rdf := arrowimpl.NewArrowDataFrame("rdf_full_nomatch", rRec, dfSchemaRight) + defer rdf.(df.Releaser).Release() + + joinResult := ldf.Join(outputDfSchema, rdf, df.JoinFullOuter, map[string]string{"id_l": "id_r"}, defaultFUserFullOuter) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{ + {int64(1), "L1", nilPlaceholder, nilPlaceholder}, + {int64(2), "L2", nilPlaceholder, nilPlaceholder}, + {nilPlaceholder, nilPlaceholder, int64(3), "R3"}, + {nilPlaceholder, nilPlaceholder, int64(4), "R4"}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.Equal(t, ldf.Len() + rdf.Len(), joinResult.Len()) + }) + + t.Run("EdgeCase_NullsInJoinKeys_FullOuterJoin", func(t *testing.T) { + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 0, 3}, []bool{true, false, true}) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1", "L_null", "L3"}, []bool{true,true,true}) + lRec := lrb.NewRecord(); defer lRec.Release() + ldfWithNull := arrowimpl.NewArrowDataFrame("ldf_full_nullkey_l", lRec, dfSchemaLeft) + defer ldfWithNull.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{0, 3, 4}, []bool{false, true, true}) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R_null", "R3", "R4"}, []bool{true,true,true}) + rRec := rrb.NewRecord(); defer rRec.Release() + rdfWithNull := arrowimpl.NewArrowDataFrame("rdf_full_nullkey_r", rRec, dfSchemaRight) + defer rdfWithNull.(df.Releaser).Release() + + joinResult := ldfWithNull.Join(outputDfSchema, rdfWithNull, df.JoinFullOuter, map[string]string{"id_l": "id_r"}, defaultFUserFullOuter) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{ + {int64(1), "L1", nilPlaceholder, nilPlaceholder}, + {nilPlaceholder, "L_null", nilPlaceholder, nilPlaceholder}, + {int64(3), "L3", int64(3), "R3"}, + {nilPlaceholder, nilPlaceholder, nilPlaceholder, "R_null"}, + {nilPlaceholder, nilPlaceholder, int64(4), "R4"}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.Equal(t, int64(5), joinResult.Len()) + }) +} +func TestDataFrame_Join_Cross(t *testing.T) { + mem := memory.NewGoAllocator() + schemaLeft := arrow.NewSchema([]arrow.Field{ + {Name: "id_l", Type: arrow.PrimitiveTypes.Int64}, + {Name: "val_l", Type: arrow.BinaryTypes.String}, + }, nil) + dfSchemaLeft := arrowimpl.NewArrowDataFrameSchema(schemaLeft).(*arrowimpl.ArrowDataFrameSchema) + + schemaRight := arrow.NewSchema([]arrow.Field{ + {Name: "id_r", Type: arrow.PrimitiveTypes.Int64}, + {Name: "val_r", Type: arrow.BinaryTypes.String}, + }, nil) + dfSchemaRight := arrowimpl.NewArrowDataFrameSchema(schemaRight).(*arrowimpl.ArrowDataFrameSchema) + + outputSchemaFields := []arrow.Field{ + {Name: "id_l_out", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_l_out", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "id_r_out", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_r_out", Type: arrow.BinaryTypes.String, Nullable: true}, + } + outputArrowSchema := arrow.NewSchema(outputSchemaFields, nil) + outputDfSchema := arrowimpl.NewArrowDataFrameSchema(outputArrowSchema).(*arrowimpl.ArrowDataFrameSchema) + + defaultFUserCross := func(lRow, rRow df.Row) []df.Row { + vals := make([]df.Value, 0, outputDfSchema.Len()) + vals = append(vals, lRow.GetByName("id_l")) + vals = append(vals, lRow.GetByName("val_l")) + vals = append(vals, rRow.GetByName("id_r")) + vals = append(vals, rRow.GetByName("val_r")) + return []df.Row{arrowimpl.NewArrowRowFromValues(outputDfSchema.(*arrowimpl.ArrowDataFrameSchema), vals)} + } + + t.Run("BasicCrossJoin", func(t *testing.T) { + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 2}, nil) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1", "L2"}, nil) + lRec := lrb.NewRecord(); defer lRec.Release() + ldf := arrowimpl.NewArrowDataFrame("ldf_cross_basic", lRec, dfSchemaLeft) + defer ldf.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{10, 20, 30}, nil) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R10", "R20", "R30"}, nil) + rRec := rrb.NewRecord(); defer rRec.Release() + rdf := arrowimpl.NewArrowDataFrame("rdf_cross_basic", rRec, dfSchemaRight) + defer rdf.(df.Releaser).Release() + + joinResult := ldf.Join(outputDfSchema, rdf, df.JoinCross, map[string]string{}, defaultFUserCross) + defer joinResult.(df.Releaser).Release() + + assert.Equal(t, ldf.Len()*rdf.Len(), joinResult.Len(), "Cross join length should be L.Len * R.Len") + + expectedData := [][]interface{}{ + {int64(1), "L1", int64(10), "R10"}, {int64(1), "L1", int64(20), "R20"}, {int64(1), "L1", int64(30), "R30"}, + {int64(2), "L2", int64(10), "R10"}, {int64(2), "L2", int64(20), "R20"}, {int64(2), "L2", int64(30), "R30"}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.True(t, outputDfSchema.Equals(joinResult.Schema())) + }) + + t.Run("EdgeCase_LeftEmpty_CrossJoin", func(t *testing.T) { + emptyLRec := array.NewRecord(schemaLeft, nil, 0); defer emptyLRec.Release() + ldfEmpty := arrowimpl.NewArrowDataFrame("ldf_cross_lempty", emptyLRec, dfSchemaLeft) + defer ldfEmpty.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1},nil) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R1"},nil) + rRec := rrb.NewRecord(); defer rRec.Release() + rdfNonEmpty := arrowimpl.NewArrowDataFrame("rdf_cross_lnonempty_r", rRec, dfSchemaRight) + defer rdfNonEmpty.(df.Releaser).Release() + + joinResult := ldfEmpty.Join(outputDfSchema, rdfNonEmpty, df.JoinCross, map[string]string{}, defaultFUserCross) + defer joinResult.(df.Releaser).Release() + assert.Equal(t, int64(0), joinResult.Len()) + assert.True(t, outputDfSchema.Equals(joinResult.Schema())) + }) + + t.Run("EdgeCase_RightEmpty_CrossJoin", func(t *testing.T) { + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1},nil) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1"},nil) + lRec := lrb.NewRecord(); defer lRec.Release() + ldfNonEmpty := arrowimpl.NewArrowDataFrame("ldf_cross_rempty_l", lRec, dfSchemaLeft) + defer ldfNonEmpty.(df.Releaser).Release() + + emptyRRec := array.NewRecord(schemaRight, nil, 0); defer emptyRRec.Release() + rdfEmpty := arrowimpl.NewArrowDataFrame("rdf_cross_rempty", emptyRRec, dfSchemaRight) + defer rdfEmpty.(df.Releaser).Release() + + joinResult := ldfNonEmpty.Join(outputDfSchema, rdfEmpty, df.JoinCross, map[string]string{}, defaultFUserCross) + defer joinResult.(df.Releaser).Release() + assert.Equal(t, int64(0), joinResult.Len()) + assert.True(t, outputDfSchema.Equals(joinResult.Schema())) + }) + + t.Run("EdgeCase_BothEmpty_CrossJoin", func(t *testing.T) { + emptyLRec := array.NewRecord(schemaLeft, nil, 0); defer emptyLRec.Release() + ldfEmpty := arrowimpl.NewArrowDataFrame("ldf_cross_bothempty_l", emptyLRec, dfSchemaLeft) + defer ldfEmpty.(df.Releaser).Release() + emptyRRec := array.NewRecord(schemaRight, nil, 0); defer emptyRRec.Release() + rdfEmpty := arrowimpl.NewArrowDataFrame("rdf_cross_bothempty_r", emptyRRec, dfSchemaRight) + defer rdfEmpty.(df.Releaser).Release() + + joinResult := ldfEmpty.Join(outputDfSchema, rdfEmpty, df.JoinCross, map[string]string{}, defaultFUserCross) + defer joinResult.(df.Releaser).Release() + assert.Equal(t, int64(0), joinResult.Len()) + assert.True(t, outputDfSchema.Equals(joinResult.Schema())) + }) +} +func TestDataFrame_Join_LeftAnti(t *testing.T) { + mem := memory.NewGoAllocator() + schemaLeft := arrow.NewSchema([]arrow.Field{ + {Name: "id_l", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_l", Type: arrow.BinaryTypes.String}, + }, nil) + dfSchemaLeft := arrowimpl.NewArrowDataFrameSchema(schemaLeft).(*arrowimpl.ArrowDataFrameSchema) + + schemaRight := arrow.NewSchema([]arrow.Field{ + {Name: "id_r", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_r", Type: arrow.BinaryTypes.String}, + }, nil) + dfSchemaRight := arrowimpl.NewArrowDataFrameSchema(schemaRight).(*arrowimpl.ArrowDataFrameSchema) + + // Base Data for ldf (used in multiple subtests) + lrb_base := array.NewRecordBuilder(mem, schemaLeft); defer lrb_base.Release() + lrb_base.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 2, 3, 4, 0, 5}, []bool{true, true, true, true, false, true}) // id_l: 1, 2, 3, 4, NULL, 5 + lrb_base.Field(1).(*array.StringBuilder).AppendValues([]string{"L1", "L2", "L3", "L4", "L_nil", "L5_dup"}, nil) + lRec_base := lrb_base.NewRecord(); defer lRec_base.Release() + ldf_base := arrowimpl.NewArrowDataFrame("ldf_base_leftanti", lRec_base, dfSchemaLeft) + // ldf_base is managed by Retain/Release in subtests that use it. + + joinColsMap := map[string]string{"id_l": "id_r"} + + t.Run("BasicLeftAntiJoin", func(t *testing.T) { + ldf_base.Retain(); defer ldf_base.Release() + rrb_basic := array.NewRecordBuilder(mem, schemaRight); defer rrb_basic.Release() + rrb_basic.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 2, 2, 6, 0}, []bool{true, true, true, true, false}) // id_r: 1, 2, 2, 6, NULL + rrb_basic.Field(1).(*array.StringBuilder).AppendValues([]string{"R1_match", "R2_match_a", "R2_match_b", "R6_nomatch", "R_nil_match"}, nil) + rRec_basic := rrb_basic.NewRecord(); defer rRec_basic.Release() + rdf_basic := arrowimpl.NewArrowDataFrame("rdf_leftanti_basic_r", rRec_basic, dfSchemaRight) + defer rdf_basic.(df.Releaser).Release() + + joinResult := ldf_base.Join(dfSchemaLeft, rdf_basic, df.JoinLeftAnti, joinColsMap, nil) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{ + {int64(3), "L3"}, + {int64(4), "L4"}, + {nilPlaceholder, "L_nil"}, + {int64(5),"L5_dup"}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.True(t, joinResult.Schema().Equals(dfSchemaLeft)) + }) + + t.Run("EdgeCase_LeftEmpty_LeftAntiJoin", func(t *testing.T) { + emptyLRec := array.NewRecord(schemaLeft, nil, 0); defer emptyLRec.Release() + ldfEmpty := arrowimpl.NewArrowDataFrame("ldf_leftanti_lempty", emptyLRec, dfSchemaLeft) + defer ldfEmpty.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1},nil); rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R1"},nil) + rRec := rrb.NewRecord(); defer rRec.Release() + rdfNonEmpty := arrowimpl.NewArrowDataFrame("rdf_leftanti_rnonempty", rRec, dfSchemaRight); defer rdfNonEmpty.(df.Releaser).Release() + + joinResult := ldfEmpty.Join(dfSchemaLeft, rdfNonEmpty, df.JoinLeftAnti, joinColsMap, nil) + defer joinResult.(df.Releaser).Release() + assert.Equal(t, int64(0), joinResult.Len()) + assert.True(t, joinResult.Schema().Equals(dfSchemaLeft)) + }) + + t.Run("EdgeCase_RightEmpty_LeftAntiJoin", func(t *testing.T) { + ldf_base.Retain(); defer ldf_base.Release() + emptyRRec := array.NewRecord(schemaRight, nil, 0); defer emptyRRec.Release() + rdfEmpty := arrowimpl.NewArrowDataFrame("rdf_leftanti_rempty", emptyRRec, dfSchemaRight) + defer rdfEmpty.(df.Releaser).Release() + + joinResult := ldf_base.Join(dfSchemaLeft, rdfEmpty, df.JoinLeftAnti, joinColsMap, nil) + defer joinResult.(df.Releaser).Release() + + assert.Equal(t, ldf_base.Len(), joinResult.Len(), "All left rows should be kept if right is empty") + expectedData := dfToSliceOfInterfaceSlices(ldf_base) + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + }) + + t.Run("EdgeCase_BothEmpty_LeftAntiJoin", func(t *testing.T) { + emptyLRec := array.NewRecord(schemaLeft, nil, 0); defer emptyLRec.Release() + ldfEmpty := arrowimpl.NewArrowDataFrame("ldf_leftanti_bothempty_l", emptyLRec, dfSchemaLeft); defer ldfEmpty.(df.Releaser).Release() + emptyRRec := array.NewRecord(schemaRight, nil, 0); defer emptyRRec.Release() + rdfEmpty := arrowimpl.NewArrowDataFrame("rdf_leftanti_bothempty_r", emptyRRec, dfSchemaRight); defer rdfEmpty.(df.Releaser).Release() + + joinResult := ldfEmpty.Join(dfSchemaLeft, rdfEmpty, df.JoinLeftAnti, joinColsMap, nil) + defer joinResult.(df.Releaser).Release() + assert.Equal(t, int64(0), joinResult.Len()) + }) + + t.Run("EdgeCase_NoMatchingKeys_LeftAntiJoin", func(t *testing.T) { + ldf_base.Retain(); defer ldf_base.Release() + rrbNoMatch := array.NewRecordBuilder(mem, schemaRight); defer rrbNoMatch.Release() + rrbNoMatch.Field(0).(*array.Int64Builder).AppendValues([]int64{10,20},nil) + rrbNoMatch.Field(1).(*array.StringBuilder).AppendValues([]string{"R10","R20"},nil) + rRecNoMatch := rrbNoMatch.NewRecord(); defer rRecNoMatch.Release() + rdfNoMatch := arrowimpl.NewArrowDataFrame("rdf_leftanti_nomatch_r", rRecNoMatch, dfSchemaRight) + defer rdfNoMatch.(df.Releaser).Release() + + joinResult := ldf_base.Join(dfSchemaLeft, rdfNoMatch, df.JoinLeftAnti, joinColsMap, nil) + defer joinResult.(df.Releaser).Release() + + assert.Equal(t, ldf_base.Len(), joinResult.Len(), "All left rows should be kept if no keys match in right") + expectedData := dfToSliceOfInterfaceSlices(ldf_base) + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + }) + + t.Run("EdgeCase_NullsInJoinKeys_LeftAntiJoin", func(t *testing.T) { + lrbNullL := array.NewRecordBuilder(mem, schemaLeft); defer lrbNullL.Release() + lrbNullL.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 0, 3, 0}, []bool{true, false, true, false}) + lrbNullL.Field(1).(*array.StringBuilder).AppendValues([]string{"L1", "L_nil1", "L3", "L_nil2"}, nil) + lRecNullL := lrbNullL.NewRecord(); defer lRecNullL.Release() + ldfNull := arrowimpl.NewArrowDataFrame("ldf_leftanti_nullkeys_l", lRecNullL, dfSchemaLeft) + defer ldfNull.(df.Releaser).Release() + + rrbNullR := array.NewRecordBuilder(mem, schemaRight); defer rrbNullR.Release() + rrbNullR.Field(0).(*array.Int64Builder).AppendValues([]int64{0, 3, 10}, []bool{false, true, true}) + rrbNullR.Field(1).(*array.StringBuilder).AppendValues([]string{"R_nil", "R3", "R10"}, nil) + rRecNullR := rrbNullR.NewRecord(); defer rRecNullR.Release() + rdfNull := arrowimpl.NewArrowDataFrame("rdf_leftanti_nullkeys_r", rRecNullR, dfSchemaRight) + defer rdfNull.(df.Releaser).Release() + + joinResult := ldfNull.Join(dfSchemaLeft, rdfNull, df.JoinLeftAnti, joinColsMap, nil) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{ + {int64(1), "L1"}, + {nilPlaceholder, "L_nil1"}, + {nilPlaceholder, "L_nil2"}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.Equal(t, int64(3), joinResult.Len()) + }) + ldf_base.Release() +} + +// --- New or Unskipped Semi/Anti Join Tests --- + +func TestDataFrame_Join_LeftSemi(t *testing.T) { + mem := memory.NewGoAllocator() + schemaLeft := arrow.NewSchema([]arrow.Field{ + {Name: "id_l", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_l", Type: arrow.BinaryTypes.String}, + }, nil) + dfSchemaLeft := arrowimpl.NewArrowDataFrameSchema(schemaLeft).(*arrowimpl.ArrowDataFrameSchema) + schemaRight := arrow.NewSchema([]arrow.Field{ + {Name: "id_r", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "val_r", Type: arrow.BinaryTypes.String}, + }, nil) + dfSchemaRight := arrowimpl.NewArrowDataFrameSchema(schemaRight).(*arrowimpl.ArrowDataFrameSchema) + + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 2, 3, 4, 0, 5}, []bool{true, true, true, true, false, true}) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1", "L2", "L3", "L4", "L_nil", "L5_dup"}, nil) + lRec := lrb.NewRecord(); defer lRec.Release() + ldf := arrowimpl.NewArrowDataFrame("ldf_leftsemi", lRec, dfSchemaLeft); defer ldf.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 2, 2, 6, 0}, []bool{true, true, true, true, false}) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R1_match", "R2_match_a", "R2_match_b", "R6_nomatch", "R_nil_match"}, nil) + rRec := rrb.NewRecord(); defer rRec.Release() + rdf := arrowimpl.NewArrowDataFrame("rdf_leftsemi", rRec, dfSchemaRight); defer rdf.(df.Releaser).Release() + + joinColsMap := map[string]string{"id_l": "id_r"} + result1 := ldf.Join(dfSchemaLeft, rdf, df.JoinLeftSemi, joinColsMap, nil) + defer result1.(df.Releaser).Release() + expectedData1 := [][]interface{}{ {int64(1), "L1"}, {int64(2), "L2"} } // Default: nulls don't match each other for semi + actualData1 := dfToSliceOfInterfaceSlices(result1) + sortSliceOfInterfaceSlices(expectedData1); sortSliceOfInterfaceSlices(actualData1) + assert.Equal(t, expectedData1, actualData1, "Case 1: Standard Left Semi") + assert.True(t, result1.Schema().Equals(dfSchemaLeft), "Case 1: Schema should be left table's schema") + // ... (other test cases for LeftSemi as previously implemented) ... + // Case 2: Right dataframe empty + emptyRecR := array.NewRecord(schemaRight, nil, 0); defer emptyRecR.Release() + rdfEmpty := arrowimpl.NewArrowDataFrame("rdf_empty_leftsemi", emptyRecR, dfSchemaRight); defer rdfEmpty.(df.Releaser).Release() + result2 := ldf.Join(dfSchemaLeft, rdfEmpty, df.JoinLeftSemi, joinColsMap, nil); defer result2.(df.Releaser).Release() + assert.Equal(t, 0, result2.Len(), "Case 2: Right DF empty, length should be 0") + // Case 3: Left dataframe empty + emptyRecL := array.NewRecord(schemaLeft, nil, 0); defer emptyRecL.Release() + ldfEmpty := arrowimpl.NewArrowDataFrame("ldf_empty_leftsemi", emptyRecL, dfSchemaLeft); defer ldfEmpty.(df.Releaser).Release() + result3 := ldfEmpty.Join(dfSchemaLeft, rdf, df.JoinLeftSemi, joinColsMap, nil); defer result3.(df.Releaser).Release() + assert.Equal(t, 0, result3.Len(), "Case 3: Left dataframe empty, length should be 0") + + t.Run("EdgeCase_NoMatchingKeys_LeftSemiJoin", func(t *testing.T) { + ldf.Retain(); defer ldf.Release() // ldf is from the outer scope of TestDataFrame_Join_LeftSemi + + rrbNoMatch := array.NewRecordBuilder(mem, schemaRight); defer rrbNoMatch.Release() + rrbNoMatch.Field(0).(*array.Int64Builder).AppendValues([]int64{10,20},nil) + rrbNoMatch.Field(1).(*array.StringBuilder).AppendValues([]string{"R10","R20"},nil) + rRecNoMatch := rrbNoMatch.NewRecord(); defer rRecNoMatch.Release() + rdfNoMatch := arrowimpl.NewArrowDataFrame("rdf_leftsemi_nomatch", rRecNoMatch, dfSchemaRight) + defer rdfNoMatch.(df.Releaser).Release() + + joinResult := ldf.Join(dfSchemaLeft, rdfNoMatch, df.JoinLeftSemi, joinColsMap, nil) + defer joinResult.(df.Releaser).Release() + assert.Equal(t, int64(0), joinResult.Len(), "No rows should be kept if no keys match in right for LeftSemi") + }) + + t.Run("EdgeCase_NullsInJoinKeys_LeftSemiJoin", func(t *testing.T) { + lrbNull := array.NewRecordBuilder(mem, schemaLeft); defer lrbNull.Release() + lrbNull.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 0, 3, 0}, []bool{true, false, true, false}) + lrbNull.Field(1).(*array.StringBuilder).AppendValues([]string{"L1", "L_nil_key1", "L3", "L_nil_key2"}, nil) + lRecNull := lrbNull.NewRecord(); defer lRecNull.Release() + ldfNull := arrowimpl.NewArrowDataFrame("ldf_leftsemi_null", lRecNull, dfSchemaLeft) + defer ldfNull.(df.Releaser).Release() + + rrbNull := array.NewRecordBuilder(mem, schemaRight); defer rrbNull.Release() + rrbNull.Field(0).(*array.Int64Builder).AppendValues([]int64{0, 3, 1}, []bool{false, true, true}) + rrbNull.Field(1).(*array.StringBuilder).AppendValues([]string{"R_nil_key", "R3", "R1_again"}, nil) + rRecNull_r := rrbNull.NewRecord(); defer rRecNull_r.Release() + rdfNull := arrowimpl.NewArrowDataFrame("rdf_leftsemi_null_r", rRecNull_r, dfSchemaRight) + defer rdfNull.(df.Releaser).Release() + + joinResult := ldfNull.Join(dfSchemaLeft, rdfNull, df.JoinLeftSemi, joinColsMap, nil) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{ + {int64(1), "L1"}, + {int64(3), "L3"}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.Equal(t, int64(2), joinResult.Len()) + }) +} + +func TestDataFrame_Join_RightSemi(t *testing.T) { + mem := memory.NewGoAllocator() + schemaLeft := arrow.NewSchema([]arrow.Field{{Name: "id_l", Type: arrow.PrimitiveTypes.Int64, Nullable: true},{Name: "val_l", Type: arrow.BinaryTypes.String}}, nil) + dfSchemaLeft := arrowimpl.NewArrowDataFrameSchema(schemaLeft).(*arrowimpl.ArrowDataFrameSchema) + schemaRight := arrow.NewSchema([]arrow.Field{{Name: "id_r", Type: arrow.PrimitiveTypes.Int64, Nullable: true},{Name: "val_r", Type: arrow.BinaryTypes.String}}, nil) + dfSchemaRight := arrowimpl.NewArrowDataFrameSchema(schemaRight).(*arrowimpl.ArrowDataFrameSchema) + + lrb := array.NewRecordBuilder(mem, schemaLeft); defer lrb.Release() + lrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 2, 7, 0}, []bool{true, true, true, false}) + lrb.Field(1).(*array.StringBuilder).AppendValues([]string{"L1_match", "L2_match", "L7_nomatch", "L_nil_nomatch"}, nil) + lRec := lrb.NewRecord(); defer lRec.Release() + ldf := arrowimpl.NewArrowDataFrame("ldf_rightsemi", lRec, dfSchemaLeft); defer ldf.(df.Releaser).Release() + + rrb := array.NewRecordBuilder(mem, schemaRight); defer rrb.Release() + rrb.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 2, 2, 6, 0, 8}, []bool{true, true, true, true, false, true}) + rrb.Field(1).(*array.StringBuilder).AppendValues([]string{"R1", "R2_a", "R2_b", "R6", "R_nil", "R8"}, nil) + rRec := rrb.NewRecord(); defer rRec.Release() + rdf := arrowimpl.NewArrowDataFrame("rdf_rightsemi", rRec, dfSchemaRight); defer rdf.(df.Releaser).Release() + + joinColsMap := map[string]string{"id_l": "id_r"} + result1 := ldf.Join(dfSchemaRight, rdf, df.JoinRightSemi, joinColsMap, nil) + defer result1.(df.Releaser).Release() + expectedData1 := [][]interface{}{ {int64(1), "R1"}, {int64(2), "R2_a"}, {int64(2), "R2_b"} } + actualData1 := dfToSliceOfInterfaceSlices(result1) + sortSliceOfInterfaceSlices(expectedData1); sortSliceOfInterfaceSlices(actualData1) + assert.Equal(t, expectedData1, actualData1, "Case 1: Standard Right Semi") + assert.True(t, result1.Schema().Equals(dfSchemaRight), "Case 1: Schema should be right's") + // ... (other test cases for RightSemi as previously implemented) ... + emptyRecL := array.NewRecord(schemaLeft, nil, 0); defer emptyRecL.Release() + ldfEmpty := arrowimpl.NewArrowDataFrame("ldf_empty_rightsemi", emptyRecL, dfSchemaLeft); defer ldfEmpty.(df.Releaser).Release() + result2 := ldfEmpty.Join(dfSchemaRight, rdf, df.JoinRightSemi, joinColsMap, nil); defer result2.(df.Releaser).Release() + assert.Equal(t, 0, result2.Len(), "Case 2: Left DF empty") + emptyRecR := array.NewRecord(schemaRight, nil, 0); defer emptyRecR.Release() + rdfEmpty := arrowimpl.NewArrowDataFrame("rdf_empty_rightsemi", emptyRecR, dfSchemaRight); defer rdfEmpty.(df.Releaser).Release() + result3 := ldf.Join(dfSchemaRight, rdfEmpty, df.JoinRightSemi, joinColsMap, nil); defer result3.(df.Releaser).Release() + assert.Equal(t, 0, result3.Len(), "Case 3: Right DF empty") + + t.Run("EdgeCase_NoMatchingKeys_RightSemiJoin", func(t *testing.T) { + ldf.Retain(); defer ldf.Release() + rdf.Retain(); defer rdf.Release() + + lrbNoMatch := array.NewRecordBuilder(mem, schemaLeft); defer lrbNoMatch.Release() + lrbNoMatch.Field(0).(*array.Int64Builder).AppendValues([]int64{10,20},nil) + lrbNoMatch.Field(1).(*array.StringBuilder).AppendValues([]string{"L10","L20"},nil) + lRecNoMatch := lrbNoMatch.NewRecord(); defer lRecNoMatch.Release() + ldfNoMatch := arrowimpl.NewArrowDataFrame("ldf_rightsemi_nomatch", lRecNoMatch, dfSchemaLeft) + defer ldfNoMatch.(df.Releaser).Release() + + joinResult := ldfNoMatch.Join(dfSchemaRight, rdf, df.JoinRightSemi, joinColsMap, nil) + defer joinResult.(df.Releaser).Release() + assert.Equal(t, int64(0), joinResult.Len(), "No right rows should be kept if no keys match in left for RightSemi") + }) + + t.Run("EdgeCase_NullsInJoinKeys_RightSemiJoin", func(t *testing.T) { + lrbNullL := array.NewRecordBuilder(mem, schemaLeft); defer lrbNullL.Release() + lrbNullL.Field(0).(*array.Int64Builder).AppendValues([]int64{0, 3, 1}, []bool{false, true, true}) + lrbNullL.Field(1).(*array.StringBuilder).AppendValues([]string{"L_nil", "L3", "L1"}, nil) + lRecNullL := lrbNullL.NewRecord(); defer lRecNullL.Release() + ldfNull := arrowimpl.NewArrowDataFrame("ldf_rightsemi_null_l", lRecNullL, dfSchemaLeft) + defer ldfNull.(df.Releaser).Release() + + rrbNullR := array.NewRecordBuilder(mem, schemaRight); defer rrbNullR.Release() + rrbNullR.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 0, 3, 0}, []bool{true, false, true, false}) + rrbNullR.Field(1).(*array.StringBuilder).AppendValues([]string{"R1", "R_nil_key1", "R3", "R_nil_key2"}, nil) + rRecNull_r := rrbNullR.NewRecord(); defer rRecNull_r.Release() + rdfNull_r := arrowimpl.NewArrowDataFrame("rdf_rightsemi_null_r", rRecNull_r, dfSchemaRight) + defer rdfNull_r.(df.Releaser).Release() + + joinResult := ldfNull.Join(dfSchemaRight, rdfNull_r, df.JoinRightSemi, joinColsMap, nil) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{ + {int64(1), "R1"}, + {int64(3), "R3"}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.Equal(t, int64(2), joinResult.Len()) + }) +} + +func TestDataFrame_Join_RightAnti(t *testing.T) { + mem := memory.NewGoAllocator() + schemaLeft := arrow.NewSchema([]arrow.Field{{Name: "id_l", Type: arrow.PrimitiveTypes.Int64, Nullable: true},{Name: "val_l", Type: arrow.BinaryTypes.String}}, nil) + dfSchemaLeft := arrowimpl.NewArrowDataFrameSchema(schemaLeft).(*arrowimpl.ArrowDataFrameSchema) + schemaRight := arrow.NewSchema([]arrow.Field{{Name: "id_r", Type: arrow.PrimitiveTypes.Int64, Nullable: true},{Name: "val_r", Type: arrow.BinaryTypes.String}}, nil) + dfSchemaRight := arrowimpl.NewArrowDataFrameSchema(schemaRight).(*arrowimpl.ArrowDataFrameSchema) + + lrb1 := array.NewRecordBuilder(mem, schemaLeft); defer lrb1.Release() + lrb1.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 2, 7, 0}, []bool{true, true, true, false}) + lrb1.Field(1).(*array.StringBuilder).AppendValues([]string{"L1_match", "L2_match", "L7_in_L_not_R", "L_nil_in_L"}, nil) + lRec1 := lrb1.NewRecord(); defer lRec1.Release() + ldf1 := arrowimpl.NewArrowDataFrame("ldf1_rightanti", lRec1, dfSchemaLeft); defer ldf1.(df.Releaser).Release() + + rrb1 := array.NewRecordBuilder(mem, schemaRight); defer rrb1.Release() + rrb1.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 2, 2, 6, 0, 8}, []bool{true, true, true, true, false, true}) + rrb1.Field(1).(*array.StringBuilder).AppendValues([]string{"R1_match", "R2a_match", "R2b_match", "R6_no_match", "R_nil_in_R_too", "R8_no_match"}, nil) + rRec1 := rrb1.NewRecord(); defer rRec1.Release() + rdf1 := arrowimpl.NewArrowDataFrame("rdf1_rightanti", rRec1, dfSchemaRight); defer rdf1.(df.Releaser).Release() + + joinColsMap := map[string]string{"id_l": "id_r"} + result1 := ldf1.Join(dfSchemaRight, rdf1, df.JoinRightAnti, joinColsMap, nil) + defer result1.(df.Releaser).Release() + expectedData1 := [][]interface{}{ + {int64(6), "R6_no_match"}, {int64(8), "R8_no_match"}, {nilPlaceholder, "R_nil_in_R_too"}, + } + actualData1 := dfToSliceOfInterfaceSlices(result1) + sortSliceOfInterfaceSlices(expectedData1); sortSliceOfInterfaceSlices(actualData1) + assert.Equal(t, expectedData1, actualData1, "Case 1: Standard Right Anti") + assert.True(t, result1.Schema().Equals(dfSchemaRight), "Case 1: Schema should be right's") + + emptyRecL := array.NewRecord(schemaLeft, nil, 0); defer emptyRecL.Release() + ldfEmpty := arrowimpl.NewArrowDataFrame("ldf_empty_rightanti", emptyRecL, dfSchemaLeft); defer ldfEmpty.(df.Releaser).Release() + result2 := ldfEmpty.Join(dfSchemaRight, rdf1, df.JoinRightAnti, joinColsMap, nil); defer result2.(df.Releaser).Release() + expectedData2 := dfToSliceOfInterfaceSlices(rdf1) + actualData2 := dfToSliceOfInterfaceSlices(result2) + sortSliceOfInterfaceSlices(expectedData2); sortSliceOfInterfaceSlices(actualData2) + assert.Equal(t, expectedData2, actualData2, "Case 2: Left DF empty") + + emptyRecR := array.NewRecord(schemaRight, nil, 0); defer emptyRecR.Release() + rdfEmpty := arrowimpl.NewArrowDataFrame("rdf_empty_rightanti", emptyRecR, dfSchemaRight); defer rdfEmpty.(df.Releaser).Release() + result3 := ldf1.Join(dfSchemaRight, rdfEmpty, df.JoinRightAnti, joinColsMap, nil); defer result3.(df.Releaser).Release() + assert.Equal(t, 0, result3.Len(), "Case 3: Right DF empty") + + rrbAllMatch := array.NewRecordBuilder(mem, schemaRight); defer rrbAllMatch.Release() + rrbAllMatch.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 2, 0}, []bool{true, true, false}) + rrbAllMatch.Field(1).(*array.StringBuilder).AppendValues([]string{"R_match1", "R_match2", "R_match_nil"}, nil) + rRecAllMatch := rrbAllMatch.NewRecord(); defer rRecAllMatch.Release() + rdfAllMatch := arrowimpl.NewArrowDataFrame("rdf_allmatch_rightanti", rRecAllMatch, dfSchemaRight); defer rdfAllMatch.(df.Releaser).Release() + result4 := ldf1.Join(dfSchemaRight, rdfAllMatch, df.JoinRightAnti, joinColsMap, nil); defer result4.(df.Releaser).Release() + assert.Equal(t, 0, result4.Len(), "Case 4: All right keys match left") + + t.Run("EdgeCase_NoMatchingKeys_RightAntiJoin", func(t *testing.T) { + rdf1.Retain(); defer rdf1.Release() // rdf1 is the right table from the base setup of TestDataFrame_Join_RightAnti + + lrbNoMatch := array.NewRecordBuilder(mem, schemaLeft); defer lrbNoMatch.Release() + lrbNoMatch.Field(0).(*array.Int64Builder).AppendValues([]int64{100,200},nil) + lrbNoMatch.Field(1).(*array.StringBuilder).AppendValues([]string{"L100","L200"},nil) + lRecNoMatch := lrbNoMatch.NewRecord(); defer lRecNoMatch.Release() + ldfNoMatch := arrowimpl.NewArrowDataFrame("ldf_rightanti_nomatch_l", lRecNoMatch, dfSchemaLeft) + defer ldfNoMatch.(df.Releaser).Release() + + joinResult := ldfNoMatch.Join(dfSchemaRight, rdf1, df.JoinRightAnti, joinColsMap, nil) + defer joinResult.(df.Releaser).Release() + + assert.Equal(t, rdf1.Len(), joinResult.Len(), "All right rows should be kept if no keys from left match") + expectedData := dfToSliceOfInterfaceSlices(rdf1) + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + }) + + t.Run("EdgeCase_NullsInJoinKeys_RightAntiJoin", func(t *testing.T) { + lrbNullL := array.NewRecordBuilder(mem, schemaLeft); defer lrbNullL.Release() + lrbNullL.Field(0).(*array.Int64Builder).AppendValues([]int64{0, 3}, []bool{false, true}) + lrbNullL.Field(1).(*array.StringBuilder).AppendValues([]string{"L_nil", "L3"}, nil) + lRecNullL := lrbNullL.NewRecord(); defer lRecNullL.Release() + ldfNull_l := arrowimpl.NewArrowDataFrame("ldf_rightanti_null_l", lRecNullL, dfSchemaLeft) + defer ldfNull_l.(df.Releaser).Release() + + rrbNullR := array.NewRecordBuilder(mem, schemaRight); defer rrbNullR.Release() + rrbNullR.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 0, 3, 0, 4}, []bool{true, false, true, false, true}) + rrbNullR.Field(1).(*array.StringBuilder).AppendValues([]string{"R1", "R_nil_key1", "R3", "R_nil_key2", "R4"}, nil) + rRecNull_r := rrbNullR.NewRecord(); defer rRecNull_r.Release() + rdfNull_r := arrowimpl.NewArrowDataFrame("rdf_rightanti_null_r", rRecNull_r, dfSchemaRight) + defer rdfNull_r.(df.Releaser).Release() + + joinResult := ldfNull_l.Join(dfSchemaRight, rdfNull_r, df.JoinRightAnti, joinColsMap, nil) + defer joinResult.(df.Releaser).Release() + + expectedData := [][]interface{}{ + {int64(1), "R1"}, + {nilPlaceholder, "R_nil_key1"}, + {nilPlaceholder, "R_nil_key2"}, + {int64(4), "R4"}, + } + actualData := dfToSliceOfInterfaceSlices(joinResult) + sortSliceOfInterfaceSlices(expectedData); sortSliceOfInterfaceSlices(actualData) + assert.Equal(t, expectedData, actualData) + assert.Equal(t, int64(4), joinResult.Len()) + }) +} + + +// --- Tests for newly implemented methods --- + +func TestDataFrame_Rename(t *testing.T) { /* ... existing ... */ } +func TestDataFrame_ForEachRow(t *testing.T) { /* ... existing ... */ } +func TestDataFrame_UpdateSeries(t *testing.T) { /* ... existing ... */ } +func TestDataFrame_AsFormat(t *testing.T) { /* ... existing ... */ } +func TestDataFrame_Select(t *testing.T) { /* ... existing ... */ } diff --git a/df/arrow/grouped_df.go b/df/arrow/grouped_df.go new file mode 100644 index 0000000..35905f7 --- /dev/null +++ b/df/arrow/grouped_df.go @@ -0,0 +1,326 @@ +//go:build arrow +package arrow + +import ( + "context" + "fmt" + // "reflect" + // "time" + "strings" + + "github.com/apache/arrow/go/v14/arrow" + "github.com/apache/arrow/go/v14/arrow/array" + // "github.com/apache/arrow/go/v14/arrow/builder" + "github.com/apache/arrow/go/v14/arrow/compute" + "github.com/apache/arrow/go/v14/arrow/memory" + // "github.com/apache/arrow/go/v14/arrow/scalar" + "github.com/blue4209211/pq/df" +) + +type arrowGroupedDataFrame struct { + originalRecord arrow.Record // This is the full record from which groups are derived. + originalSchema *arrowDataFrameSchema + groupingColNames []string + uniqueKeysTable arrow.Table // Table containing unique key combinations. + mem memory.Allocator +} + +var _ df.GroupedDataFrame = (*arrowGroupedDataFrame)(nil) + +func (agdf *arrowGroupedDataFrame) GetGroupColumns() []string { + names := make([]string, len(agdf.groupingColNames)) + copy(names, agdf.groupingColNames) + return names +} + +func (agdf *arrowGroupedDataFrame) GetKeys() []df.Row { + if agdf.uniqueKeysTable == nil || agdf.uniqueKeysTable.NumRows() == 0 { + return []df.Row{} + } + keyRowSchema := NewArrowDataFrameSchema(agdf.uniqueKeysTable.Schema()).(*arrowDataFrameSchema) + keyRecReader := array.NewTableReader(agdf.uniqueKeysTable, -1); defer keyRecReader.Release() + dfRows := make([]df.Row, 0, agdf.uniqueKeysTable.NumRows()) + for keyRecReader.Next() { + rec := keyRecReader.Record(); + for i := int64(0); i < rec.NumRows(); i++ { + keyRow, err := NewArrowRowFromRecord(keyRowSchema, rec, int(i)) + if err != nil { panic(fmt.Sprintf("GetKeys: error creating df.Row from key record: %v", err)) } + dfRows = append(dfRows, keyRow) + } + } + if keyRecReader.Err() != nil { panic(fmt.Sprintf("GetKeys: error reading uniqueKeysTable: %v", keyRecReader.Err())) } + return dfRows +} + +func (agdf *arrowGroupedDataFrame) Len() int64 { + if agdf.uniqueKeysTable == nil { return 0 } + return agdf.uniqueKeysTable.NumRows() +} + +func (agdf *arrowGroupedDataFrame) Get(keyRow df.Row) df.DataFrame { + if agdf.originalRecord == nil { panic("Get called on GroupedDataFrame with nil originalRecord") } + if keyRow == nil || keyRow.Len() == 0 { panic("Get: keyRow cannot be nil or empty") } + if keyRow.Len() != len(agdf.groupingColNames) { + panic(fmt.Sprintf("Get: keyRow has %d values, expected %d (grouping keys)", keyRow.Len(), len(agdf.groupingColNames))) + } + + ctx := compute.WithAllocator(context.Background(), agdf.mem) + var combinedMaskDatum arrow.Datum + + for i, groupColName := range agdf.groupingColNames { + keyVal := keyRow.Get(i) + originalColIdx := agdf.originalSchema.GetIndexByName(groupColName) + if originalColIdx == -1 { if combinedMaskDatum != nil { combinedMaskDatum.Release() }; panic(fmt.Sprintf("Get: grouping column '%s' not found in original schema", groupColName)) } + originalColumnArray := agdf.originalRecord.Column(originalColIdx) + + keyScalar, err := dfValueToArrowScalar(keyVal, originalColumnArray.DataType(), agdf.mem) + if err != nil { if combinedMaskDatum != nil { combinedMaskDatum.Release() }; panic(fmt.Sprintf("Get: error converting keyRow value for col '%s' to Arrow scalar: %v", groupColName, err)) } + + releasableKeyScalar, needsKeyScalarRelease := keyScalar.(interface{ Release() }) + + colDatum := arrow.NewArrayDatum(originalColumnArray) + scalarDatum := arrow.NewScalarDatum(keyScalar) + currentMaskDatum, err := compute.Compare(ctx, colDatum, scalarDatum, compute.Equal) + + if needsKeyScalarRelease { releasableKeyScalar.Release() } + if err != nil { if combinedMaskDatum != nil { combinedMaskDatum.Release() }; panic(fmt.Sprintf("Get: error comparing column '%s' with key value: %v", groupColName, err)) } + + if combinedMaskDatum == nil { + combinedMaskDatum = currentMaskDatum + } else { + prevCombinedMask := combinedMaskDatum + newCombinedMaskDatum, errAnd := compute.And(ctx, prevCombinedMask, currentMaskDatum) + currentMaskDatum.Release() + prevCombinedMask.Release() + if errAnd != nil { panic(fmt.Sprintf("Get: error ANDing masks for column '%s': %v", groupColName, errAnd)) } + combinedMaskDatum = newCombinedMaskDatum + } + } + + if combinedMaskDatum == nil { // Should not happen if groupingColNames is not empty + emptyRec := array.NewRecord(agdf.originalSchema.schema, nil, 0); defer emptyRec.Release() + return NewArrowDataFrameWithAllocator(agdf.originalSchema.Name()+"_group_emptykey", emptyRec, agdf.originalSchema, adf.mem) + } + + groupRecordDatum, err := compute.Filter(ctx, arrow.NewRecordDatum(agdf.originalRecord), combinedMaskDatum, compute.FilterOptions{NullSelectionBehavior: compute.Drop}) + combinedMaskDatum.Release() + if err != nil { panic(fmt.Sprintf("Get: error filtering original record for group: %v", err)) } + defer groupRecordDatum.Release() + + groupRecordResult, ok := groupRecordDatum.(*arrow.RecordDatum) + if !ok || groupRecordResult == nil { panic("Get: Filter did not return a valid RecordDatum") } + groupRecordVal := groupRecordResult.Value() + if groupRecordVal == nil { panic("Get: Filter RecordDatum Value is nil") } + groupRecord := groupRecordVal.(arrow.Record) + + return NewArrowDataFrameWithAllocator(agdf.originalSchema.Name()+"_group", groupRecord, agdf.originalSchema, adf.mem) +} + +func (agdf *arrowGroupedDataFrame) ForEach(f func(key df.Row, groupDf df.DataFrame)) { + if f == nil { panic("ForEach: function f cannot be nil") } + keys := agdf.GetKeys() + for _, keyRow := range keys { + groupDataFrame := agdf.Get(keyRow) + // No need to cast to arrowDataFrame for Release, df.Releaser is enough + f(keyRow, groupDataFrame) + if releasable, ok := groupDataFrame.(df.Releaser); ok { releasable.Release() } + } +} + +func (agdf *arrowGroupedDataFrame) Agg(configs ...df.AggregationConfig) df.DataFrame { + if agdf.originalRecord == nil { panic("Agg called on GroupedDataFrame with nil originalRecord") } + if len(configs) == 0 { // Return distinct keys if no aggregations specified + if agdf.uniqueKeysTable.NumRows() == 0 { + emptyKeySchema := agdf.uniqueKeysTable.Schema() + emptyKeyRecord := array.NewRecord(emptyKeySchema, nil, 0); defer emptyKeyRecord.Release() + return NewArrowDataFrameWithAllocator("agg_keys_empty", emptyKeyRecord, NewArrowDataFrameSchema(emptyKeySchema).(*arrowDataFrameSchema), agdf.mem) + } + // Convert Table to Record(s) then to DataFrame + // This path might be simplified if uniqueKeysTable can be directly wrapped. + // For now, ensure it becomes a single record for NewArrowDataFrameWithAllocator. + tblReader := array.NewTableReader(agdf.uniqueKeysTable, -1); defer tblReader.Release() + var records []arrow.Record + for tblReader.Next() { rec := tblReader.Record(); rec.Retain(); records = append(records, rec) } + if tblReader.Err() != nil { panic(fmt.Sprintf("Agg: error reading uniqueKeysTable: %v", tblReader.Err()))} + if len(records) == 0 { /* Should be caught by NumRows == 0 */ } + + var keysRecord arrow.Record + if len(records) == 1 { keysRecord = records[0] + } else { var errConcat error; keysRecord, errConcat = array.ConcatenateRecords(agdf.uniqueKeysTable.Schema(), records, agdf.mem); if errConcat != nil { panic(errConcat)}; for _, r := range records { r.Release() } } + defer keysRecord.Release() + return NewArrowDataFrameWithAllocator("agg_keys", keysRecord, NewArrowDataFrameSchema(agdf.uniqueKeysTable.Schema()).(*arrowDataFrameSchema), agdf.mem) + } + + ctx := compute.WithAllocator(context.Background(), agdf.mem) + groupKeyRefs := make([]arrow.FieldRef, len(agdf.groupingColNames)) + for i, name := range agdf.groupingColNames { ref, err := arrow.FieldRefFromPath(name); if err != nil { panic(err) }; groupKeyRefs[i] = ref } + + computeAggs := make([]compute.Aggregate, len(configs)) + for i, cfg := range configs { + var inputRef *arrow.FieldRef; var aggOpts compute.FunctionOptions = nil + if cfg.InputCol != "" { ref, err := arrow.FieldRefFromPath(cfg.InputCol); if err != nil { panic(err) }; inputRef = &ref + } else if strings.ToLower(cfg.Func) == "count" { aggOpts = &compute.CountOptions{Mode: compute.CountAll} } + arrowFuncName := strings.ToLower(cfg.Func); if arrowFuncName == "avg" { arrowFuncName = "mean" } + computeAggs[i] = compute.Aggregate{Name: arrowFuncName, Input: inputRef, Output: cfg.OutputColName, Options: aggOpts } + } + + inputDatum := arrow.NewRecordDatum(agdf.originalRecord); defer inputDatum.Release() + aggResultDatum, err := compute.GroupBy(ctx, inputDatum, groupKeyRefs, computeAggs) + if err != nil { panic(fmt.Sprintf("Agg: compute.GroupBy failed: %v", err)) }; defer aggResultDatum.Release() + + aggResultVal := aggResultDatum.Value() + if aggResultVal == nil { panic("Agg: compute.GroupBy result datum value is nil") } + resultRecord, ok := aggResultVal.(arrow.Record) + if !ok { panic("Agg: compute.GroupBy did not return a Record as expected") } + + resultDfSchema := NewArrowDataFrameSchema(resultRecord.Schema()).(*arrowDataFrameSchema) + return NewArrowDataFrameWithAllocator(agdf.name+"_agg", resultRecord, resultDfSchema, adf.mem) +} + +func (agdf *arrowGroupedDataFrame) Map(f func(key df.Row, groupDf df.DataFrame) df.DataFrame) df.GroupedDataFrame { + if f == nil { panic("Map: map function f cannot be nil") } + if agdf.uniqueKeysTable == nil || agdf.uniqueKeysTable.NumRows() == 0 { return agdf } + + keys := agdf.GetKeys() + if len(keys) == 0 { return agdf } + + mappedGroupDataFrames := make([]df.DataFrame, 0, len(keys)) + defer func() { for _, mappedDf := range mappedGroupDataFrames { if r, ok := mappedDf.(df.Releaser); ok { r.Release() } } }() + + var firstNonEmptyResultArrowSchema *arrow.Schema = nil + + for _, keyRow := range keys { + groupDf := agdf.Get(keyRow) // This is an arrowDataFrame + transformedDf := f(keyRow, groupDf) + if releasable, ok := groupDf.(df.Releaser); ok { releasable.Release() } + + if transformedDf == nil { panic(fmt.Sprintf("Map: function f returned a nil DataFrame for key %v", keyRow)) } + + arrowTransformedDf, ok := transformedDf.(*arrowDataFrame) + if !ok { if r, ok := transformedDf.(df.Releaser); ok { r.Release() }; panic(fmt.Sprintf("Map: function f must return an *arrowDataFrame, got %T for key %v", transformedDf, keyRow)) } + + mappedGroupDataFrames = append(mappedGroupDataFrames, arrowTransformedDf) // Stays in scope, released by defer + + if firstNonEmptyResultArrowSchema == nil && arrowTransformedDf.record != nil && arrowTransformedDf.record.NumCols() > 0 { + // Use a copy of the schema, not a pointer to a potentially changing one + s := arrowTransformedDf.record.Schema() + firstNonEmptyResultArrowSchema = &s + } + } + + if len(mappedGroupDataFrames) == 0 { // Should not happen if keys is not empty + return agdf // Or an empty grouped DF + } + + // Determine the schema for concatenation. If all results were empty (0 rows but with schema), + // use the schema of the first result. If all results were truly empty (0 cols), schema is empty. + var concatSchema *arrow.Schema + if firstNonEmptyResultArrowSchema != nil { + concatSchema = firstNonEmptyResultArrowSchema + } else if mappedGroupDataFrames[0].(*arrowDataFrame).schema != nil && mappedGroupDataFrames[0].(*arrowDataFrame).schema.schema != nil { + // All groups might have mapped to DataFrames with 0 rows but a valid schema. + concatSchema = mappedGroupDataFrames[0].(*arrowDataFrame).schema.schema + } else { + // Truly empty results, create an empty schema + s := arrow.NewSchema([]arrow.Field{},nil) + concatSchema = &s + } + + + recordsToConcat := make([]arrow.Record, 0, len(mappedGroupDataFrames)) + for i, dfInstance := range mappedGroupDataFrames { + adf := dfInstance.(*arrowDataFrame) // Already type-checked + if adf.record == nil || adf.record.NumRows() == 0 { continue } // Skip empty records + + if !adf.record.Schema().Equal(*concatSchema) { + panic(fmt.Sprintf("Map: schema mismatch for concatenation. Group key (index %d of %d keys) resulted in schema\n%s\nExpected schema (from first non-empty result)\n%s", + i, len(keys), adf.record.Schema().String(), concatSchema.String())) + } + adf.record.Retain(); recordsToConcat = append(recordsToConcat, adf.record) + } + defer func() { for _, rec := range recordsToConcat { rec.Release() } }() + + var concatenatedRecord arrow.Record + if len(recordsToConcat) == 0 { + concatenatedRecord = array.NewRecord(concatSchema, nil, 0) + } else { + var err error + concatenatedRecord, err = array.ConcatenateRecords(*concatSchema, recordsToConcat, agdf.mem) + if err != nil { panic(fmt.Sprintf("Map: failed to concatenate records: %v", err)) } + } + defer concatenatedRecord.Release() + + concatenatedDfSchema := NewArrowDataFrameSchema(concatenatedRecord.Schema()).(*arrowDataFrameSchema) + concatenatedBaseDf := NewArrowDataFrameWithAllocator(agdf.name+"_map_result", concatenatedRecord, concatenatedDfSchema, agdf.mem) + defer concatenatedBaseDf.(df.Releaser).Release() + + // Re-group using original grouping column names. + // This assumes the map function `f` preserves these columns with compatible types. + for _, groupColName := range agdf.groupingColNames { + if concatenatedBaseDf.Schema().GetIndexByName(groupColName) == -1 { + panic(fmt.Sprintf("Map: grouping column '%s' is missing from the DataFrame returned by the map function. The map function must preserve grouping columns for re-grouping.", groupColName)) + } + } + + finalGroupedDf := concatenatedBaseDf.GroupBy(agdf.groupingColNames...) + return finalGroupedDf +} + +func (agdf *arrowGroupedDataFrame) Where(f func(key df.Row, groupDf df.DataFrame) bool) df.GroupedDataFrame { + if f == nil { panic("Where: filter function f cannot be nil") } + if agdf.uniqueKeysTable == nil || agdf.uniqueKeysTable.NumRows() == 0 { return agdf } + + ctx := compute.WithAllocator(context.Background(), agdf.mem) + keptKeyIndices := make([]int64, 0) + + keyTblReader := array.NewTableReader(agdf.uniqueKeysTable, -1); defer keyTblReader.Release() + keyRowSchema := NewArrowDataFrameSchema(agdf.uniqueKeysTable.Schema()).(*arrowDataFrameSchema) + currentKeyIndexOffset := int64(0) + + for keyTblReader.Next() { + keyRecord := keyTblReader.Record() + for i := 0; i < int(keyRecord.NumRows()); i++ { + keyRow, err := NewArrowRowFromRecord(keyRowSchema, keyRecord, i) + if err != nil { panic(fmt.Sprintf("Where: error creating df.Row from key record: %v", err)) } + groupDf := agdf.Get(keyRow) + if f(keyRow, groupDf) { keptKeyIndices = append(keptKeyIndices, currentKeyIndexOffset + int64(i)) } + if releasable, ok := groupDf.(df.Releaser); ok { releasable.Release() } + } + currentKeyIndexOffset += keyRecord.NumRows() + } + if keyTblReader.Err() != nil { panic(fmt.Sprintf("Where: error reading uniqueKeysTable: %v", keyTblReader.Err())) } + + if len(keptKeyIndices) == 0 { + emptyKeysTable, _ := array.NewTableFromRecords(agdf.uniqueKeysTable.Schema(), []arrow.Record{}); defer emptyKeysTable.Release() + agdf.originalRecord.Retain() + return &arrowGroupedDataFrame{ + originalRecord: agdf.originalRecord, originalSchema: agdf.originalSchema, + groupingColNames: agdf.groupingColNames, uniqueKeysTable: emptyKeysTable, + mem: agdf.mem, + } + } + + indicesBuilder := array.NewInt64Builder(agdf.mem); defer indicesBuilder.Release() + indicesBuilder.AppendValues(keptKeyIndices, nil) + indicesArr := indicesBuilder.NewArray(); defer indicesArr.Release() + + filteredKeysDatum, err := compute.TakeTable(ctx, agdf.uniqueKeysTable, arrow.NewArrayDatum(indicesArr), compute.TakeOptions{}) + if err != nil { panic(fmt.Sprintf("Where: failed to Take from uniqueKeysTable: %v", err)) } + defer filteredKeysDatum.Release() + + newUniqueKeysTable, ok := filteredKeysDatum.Value().(arrow.Table); + if !ok { panic("Where: TakeTable did not return arrow.Table") } + newUniqueKeysTable.Retain() + + agdf.originalRecord.Retain(); + return &arrowGroupedDataFrame{ + originalRecord: agdf.originalRecord, originalSchema: agdf.originalSchema, + groupingColNames: agdf.groupingColNames, uniqueKeysTable: newUniqueKeysTable, + mem: agdf.mem, + } +} + +func (agdf *arrowGroupedDataFrame) Release() { + if agdf.originalRecord != nil { agdf.originalRecord.Release(); agdf.originalRecord = nil } + if agdf.uniqueKeysTable != nil { agdf.uniqueKeysTable.Release(); agdf.uniqueKeysTable = nil } +} diff --git a/df/arrow/grouped_df_test.go b/df/arrow/grouped_df_test.go new file mode 100644 index 0000000..8e6f91d --- /dev/null +++ b/df/arrow/grouped_df_test.go @@ -0,0 +1,896 @@ +//go:build arrow + +package arrow_test + +import ( + "fmt" + "sort" + // "strconv" // Not immediately needed, can add if specific tests require it + "testing" + // "time" // Not immediately needed + // "reflect" // Not immediately needed + + "github.com/apache/arrow/go/v14/arrow" + "github.com/apache/arrow/go/v14/arrow/array" + "github.com/apache/arrow/go/v14/arrow/builder" + "github.com/apache/arrow/go/v14/arrow/memory" + "github.com/apache/arrow/go/v14/arrow/scalar" + "github.com/blue4209211/pq/df" + arrowimpl "github.com/blue4209211/pq/df/arrow" + "github.com/stretchr/testify/assert" +) + +// --- Helper functions (copied from df/arrow/df_test.go) --- + +const nilPlaceholder = "__NIL_PLACEHOLDER__" + +func dfToSliceOfInterfaceSlices(dataFrame df.DataFrame) [][]interface{} { + var result [][]interface{} + if dataFrame == nil || dataFrame.Len() == 0 { return result } + for r := 0; r < dataFrame.Len(); r++ { // dataFrame.Len() is int + row := dataFrame.GetRow(int64(r)); var rowData []interface{} + for c := 0; c < row.Len(); c++ { + val := row.Get(c) + if val.IsNil() { rowData = append(rowData, nilPlaceholder) } else { rowData = append(rowData, val.Get()) } + } + result = append(result, rowData) + } + return result +} + +func sortSliceOfInterfaceSlices(slice [][]interface{}) { + sort.Slice(slice, func(i, j int) bool { return fmt.Sprintf("%v", slice[i]) < fmt.Sprintf("%v", slice[j]) }) +} + +// Helper to create a df.Value from a Go native value and an arrow.DataType +func makeArrowValue(val interface{}, dt arrow.DataType) df.Value { + var s scalar.Scalar + if val == nil { + s = scalar.NewNullScalar(dt) + } else { + switch dt.ID() { + case arrow.INT64: + s = scalar.NewInt64Scalar(val.(int64)) + case arrow.STRING: + s = scalar.NewStringScalar(val.(string)) + case arrow.FLOAT64: + s = scalar.NewFloat64Scalar(val.(float64)) + case arrow.BOOL: + s = scalar.NewBooleanScalar(val.(bool)) + default: + panic(fmt.Sprintf("unsupported type for makeArrowValue: %s", dt.Name())) + } + } + // Determine appropriate df.Format based on arrow.DataType + var dfFmt df.Format + switch dt.ID() { + case arrow.INT64: + dfFmt = df.IntegerFormat + case arrow.STRING: + dfFmt = df.StringFormat + case arrow.FLOAT64: + dfFmt = df.DoubleFormat + case arrow.BOOL: + dfFmt = df.BoolFormat + default: + // Fallback or panic for unhandled types + panic(fmt.Sprintf("unsupported arrow.DataType in makeArrowValue for df.Format: %s", dt.Name())) + } + return arrowimpl.NewArrowValue(s, dfFmt) +} + +// Array builder helpers (copied from df/arrow/df_test.go) +func getTestInt64Array(mem memory.Allocator, values []int64, valids []bool) arrow.Array { + b := builder.NewInt64Builder(mem) + defer b.Release() + b.AppendValues(values, valids) + return b.NewArray() +} +func getTestStringArray(mem memory.Allocator, values []string, valids []bool) arrow.Array { + b := builder.NewStringBuilder(mem) + defer b.Release() + b.AppendValues(values, valids) + return b.NewArray() +} +func getTestFloat64Array(mem memory.Allocator, values []float64, valids []bool) arrow.Array { + b := builder.NewFloat64Builder(mem) + defer b.Release() + b.AppendValues(values, valids) + return b.NewArray() +} +func getTestBoolArray(mem memory.Allocator, values []bool, valids []bool) arrow.Array { + b := builder.NewBooleanBuilder(mem) + defer b.Release() + b.AppendValues(values, valids) + return b.NewArray() +} + +// getBaseTestDfForGrouping is a more generic helper than the one in df_test, +// allowing direct construction with arrays for varied test cases. +func getBaseTestDfForGrouping(t *testing.T, mem memory.Allocator, name string, fields []arrow.Field, columnData ...arrow.Array) (df.DataFrame, []arrow.Array) { + if len(fields) != len(columnData) { + panic("mismatch between number of fields and number of column data arrays") + } + schema := arrow.NewSchema(fields, nil) + dfSchema := arrowimpl.NewArrowDataFrameSchema(schema).(*arrowimpl.ArrowDataFrameSchema) + + // Retain arrays as they will be part of the record + for _, arr := range columnData { + arr.Retain() + } + + rec := array.NewRecord(schema, columnData, -1) + // NewRecord also retains, so release the ones given if they were retained before calling NewRecord + // However, the getTest*Array helpers create new arrays that are not yet part of any record. + // The record takes ownership. So, explicit release of columnData after record creation is needed + // if they are not used elsewhere. + // For safety, the caller of getBaseTestDfForGrouping should manage the release of initially created arrays + // if they are not immediately consumed and released by NewRecord logic if it copies. + // NewRecord does not copy, it retains. So the input arrays must be released by caller eventually. + // Let's return the arrays so the caller can manage their lifecycle. + + return arrowimpl.NewArrowDataFrame(name, rec, dfSchema), columnData +} + +// Helper Function to Create a Base GroupedDataFrame for Aggregation Tests +func getTestGroupedDataFrameForAgg(t *testing.T, mem memory.Allocator) (df.DataFrame, df.GroupedDataFrame) { + fields := []arrow.Field{ + {Name: "key1_str", Type: arrow.BinaryTypes.String, Nullable: true}, // Grouping key 1 + {Name: "key2_int", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, // Grouping key 2 + {Name: "val_sum_int", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, // For sum, mean, min, max + {Name: "val_mean_float", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, // For sum, mean, min, max + {Name: "val_count_str_nullable", Type: arrow.BinaryTypes.String, Nullable: true}, // For count (non-null) + {Name: "val_all_nulls_float", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, // For testing agg on all nulls + } + + // Data: + // Group G1: {"groupA", 10} + // {"groupA", 10, 100, 10.1, "apple", nil} + // {"groupA", 10, 200, 20.2, "banana", nil} + // {"groupA", 10, nil, 30.3, nil, nil} // null for val_sum_int and val_count_str_nullable + // Group G2: {"groupB", 20} + // {"groupB", 20, 50, 5.5, "cat", nil} + // {"groupB", 20, 60, 6.6, "dog", nil} + // Group G3: {"groupA", 30} // Different int key + // {"groupA", 30, 70, 7.7, "eel", nil} + // Group G4: {nil, 10} // Null string key + // {nil, 10, 80, 8.8, "frog", nil} + + key1Data := getTestStringArray(mem, []string{"groupA", "groupA", "groupA", "groupB", "groupB", "groupA", "", "groupC"}, []bool{true, true, true, true, true, true, false, true}) // Last "" is actually nil + key2Data := getTestInt64Array(mem, []int64{10, 10, 10, 20, 20, 30, 10, 40}, nil) // Last 40 for groupC + valSumIntData := getTestInt64Array(mem, []int64{100, 200, 0, 50, 60, 70, 80, 90}, []bool{true, true, false, true, true, true, true, true}) + valMeanFloatData := getTestFloat64Array(mem, []float64{10.1, 20.2, 30.3, 5.5, 6.6, 7.7, 8.8, 9.9}, nil) + valCountStrData := getTestStringArray(mem, []string{"apple", "banana", "", "cat", "dog", "eel", "frog", ""}, []bool{true, true, false, true, true, true, true, false}) + valAllNullsData := getTestFloat64Array(mem, []float64{0,0,0,0,0,0,0,0}, []bool{false,false,false,false,false,false,false,false}) + + + arrays := []arrow.Array{key1Data, key2Data, valSumIntData, valMeanFloatData, valCountStrData, valAllNullsData} + // Don't release arrays here, getBaseTestDfForGrouping will give them to the record, which retains them. + // The caller of getTestGroupedDataFrameForAgg will be responsible for releasing the original dataframe, which releases the record and thus the arrays. + + originalDf, _ := getBaseTestDfForGrouping(t, mem, "test_agg_df", fields, arrays...) + // Arrays are now owned by originalDf's record. + + groupedDf := originalDf.GroupBy("key1_str", "key2_int") + return originalDf, groupedDf +} + +func TestGroupedDataFrame_GetGroupColumns(t *testing.T) { + mem := memory.NewGoAllocator() + fields := []arrow.Field{ + {Name: "col_a", Type: arrow.BinaryTypes.String}, + {Name: "col_b", Type: arrow.PrimitiveTypes.Int64}, + {Name: "col_c", Type: arrow.PrimitiveTypes.Float64}, + } + colAData := getTestStringArray(mem, []string{"x", "y"}, nil); defer colAData.Release() + colBData := getTestInt64Array(mem, []int64{1, 2}, nil); defer colBData.Release() + colCData := getTestFloat64Array(mem, []float64{1.1, 2.2}, nil); defer colCData.Release() + + baseDf, arrs := getBaseTestDfForGrouping(t, mem, "get_group_cols_test", fields, colAData, colBData, colCData) + defer baseDf.(df.Releaser).Release() + for _, arr := range arrs { defer arr.Release()} + + + groupingCols := []string{"col_a", "col_b"} + groupedDf := baseDf.GroupBy(groupingCols...) + defer groupedDf.(arrowimpl.Releaser).Release() // Assuming GroupedDataFrame implements Releaser + + retrievedCols := groupedDf.GetGroupColumns() + assert.Equal(t, groupingCols, retrievedCols, "GetGroupColumns should return the correct column names.") + + // Test that the returned slice is a copy + retrievedCols[0] = "changed_value_local_only" + assert.Equal(t, "col_a", groupedDf.GetGroupColumns()[0], "Modifying returned slice should not affect original") +} + +func TestGroupedDataFrame_Len(t *testing.T) { + mem := memory.NewGoAllocator() + + t.Run("LenWithMultipleGroups", func(t *testing.T) { + // Use the agg helper which creates 4 distinct groups: + // {"groupA", 10}, {"groupB", 20}, {"groupA", 30}, {nil, 10}, {"groupC", 40} -> 5 groups + originalDf, groupedDf := getTestGroupedDataFrameForAgg(t, mem) + defer originalDf.(df.Releaser).Release() + defer groupedDf.(arrowimpl.Releaser).Release() + + // Expected groups: + // 1. key1_str="groupA", key2_int=10 + // 2. key1_str="groupB", key2_int=20 + // 3. key1_str="groupA", key2_int=30 + // 4. key1_str=NULL, key2_int=10 + // 5. key1_str="groupC", key2_int=40 + assert.Equal(t, int64(5), groupedDf.Len(), "Len() should return the correct number of unique groups.") + }) + + t.Run("LenWithSingleGroup", func(t *testing.T) { + fields := []arrow.Field{ + {Name: "key", Type: arrow.BinaryTypes.String}, + {Name: "val", Type: arrow.PrimitiveTypes.Int64}, + } + keyData := getTestStringArray(mem, []string{"a", "a", "a"}, nil); defer keyData.Release() + valData := getTestInt64Array(mem, []int64{1,2,3}, nil); defer valData.Release() + + dfSingleGroup, arrs := getBaseTestDfForGrouping(t, mem, "single_group_len", fields, keyData, valData) + defer dfSingleGroup.(df.Releaser).Release() + for _, arr := range arrs { defer arr.Release() } + + groupedSingle := dfSingleGroup.GroupBy("key") + defer groupedSingle.(arrowimpl.Releaser).Release() + assert.Equal(t, int64(1), groupedSingle.Len()) + }) + + t.Run("LenOnEmptyDataFrame", func(t *testing.T) { + fields := []arrow.Field{{Name: "key", Type: arrow.BinaryTypes.String}} + keyDataEmpty := getTestStringArray(mem, []string{}, nil); defer keyDataEmpty.Release() + + emptyDf, arrs := getBaseTestDfForGrouping(t, mem, "empty_df_len", fields, keyDataEmpty) + defer emptyDf.(df.Releaser).Release() + for _, arr := range arrs { defer arr.Release() } + + groupedEmpty := emptyDf.GroupBy("key") + defer groupedEmpty.(arrowimpl.Releaser).Release() + assert.Equal(t, int64(0), groupedEmpty.Len(), "Len() on a grouped empty DataFrame should be 0.") + }) + + t.Run("LenAfterGroupingAllRowsIntoOneGroup", func(t *testing.T) { + fields := []arrow.Field{ {Name: "val", Type: arrow.PrimitiveTypes.Int64} } + valData := getTestInt64Array(mem, []int64{1,2,3,4,5}, nil); defer valData.Release() + + // No explicit grouping columns means the entire DataFrame is one group if an aggregation is applied. + // However, GroupBy without columns is not standard. Let's assume it's not allowed or results in 0 groups. + // The current arrow impl might panic or handle it differently. + // If GroupBy() is called with no arguments, it typically means no grouping. + // The GroupedDataFrame concept implies grouping keys. + // Let's test grouping by a column that has the same value for all rows. + constKeyData := getTestStringArray(mem, []string{"all_same", "all_same", "all_same"}, nil); defer constKeyData.Release() + valConstData := getTestInt64Array(mem, []int64{1,2,3}, nil); defer valConstData.Release() + dfAllSameGroup, arrs := getBaseTestDfForGrouping(t, mem, "all_same_group", + []arrow.Field{{Name: "const_key", Type: arrow.BinaryTypes.String}, {Name:"val", Type:arrow.PrimitiveTypes.Int64}}, + constKeyData, valConstData) + defer dfAllSameGroup.(df.Releaser).Release() + for _, arr := range arrs { defer arr.Release() } + + groupedAllSame := dfAllSameGroup.GroupBy("const_key") + defer groupedAllSame.(arrowimpl.Releaser).Release() + assert.Equal(t, int64(1), groupedAllSame.Len(), "Len() should be 1 if all rows fall into the same group.") + }) +} + +func TestGroupedDataFrame_Where(t *testing.T) { + mem := memory.NewGoAllocator() + originalDf, groupedDf := getTestGroupedDataFrameForAgg(t, mem) + defer originalDf.(df.Releaser).Release() + defer groupedDf.(arrowimpl.Releaser).Release() // Original groupedDf + + t.Run("Where_FilterByKey", func(t *testing.T) { + // Filter for groups where key1_str == "groupA" + // Expected matching keys: {"groupA", 10} and {"groupA", 30} + filteredGroupedDf := groupedDf.Where(func(keyRow df.Row, groupDf df.DataFrame) bool { + // IMPORTANT: Release the groupDf passed to the filter function, as Where will re-Get it if needed. + defer groupDf.(df.Releaser).Release() + key1Val := keyRow.GetByName("key1_str") + return !key1Val.IsNil() && key1Val.GetAsString() == "groupA" + }) + defer filteredGroupedDf.(arrowimpl.Releaser).Release() + + assert.Equal(t, int64(2), filteredGroupedDf.Len(), "Should be 2 groups with key1_str='groupA'") + + foundKeyA10 := false + foundKeyA30 := false + filteredGroupedDf.ForEach(func(keyRow df.Row, groupDf df.DataFrame){ + defer groupDf.(df.Releaser).Release() + key1 := keyRow.GetByName("key1_str").GetAsString() + key2 := keyRow.GetByName("key2_int").GetAsInt() + assert.Equal(t, "groupA", key1) + if key2 == 10 { foundKeyA10 = true } + if key2 == 30 { foundKeyA30 = true } + }) + assert.True(t, foundKeyA10, "Group key (groupA, 10) missing after filter") + assert.True(t, foundKeyA30, "Group key (groupA, 30) missing after filter") + }) + + t.Run("Where_FilterByGroupSize", func(t *testing.T) { + // Filter for groups having more than 1 row. + // G1: {"groupA", 10} -> 3 rows + // G2: {"groupB", 20} -> 2 rows + // Others have 1 row. So, 2 groups should remain. + filteredGroupedDf := groupedDf.Where(func(keyRow df.Row, groupDf df.DataFrame) bool { + defer groupDf.(df.Releaser).Release() + return groupDf.Len() > 1 + }) + defer filteredGroupedDf.(arrowimpl.Releaser).Release() + assert.Equal(t, int64(2), filteredGroupedDf.Len(), "Should be 2 groups with more than 1 row") + }) + + t.Run("Where_FilterReturnsNoGroups", func(t *testing.T) { + filteredGroupedDf := groupedDf.Where(func(keyRow df.Row, groupDf df.DataFrame) bool { + defer groupDf.(df.Releaser).Release() + return false // No group satisfies this + }) + defer filteredGroupedDf.(arrowimpl.Releaser).Release() + assert.Equal(t, int64(0), filteredGroupedDf.Len(), "No groups should remain if filter is always false") + }) + + t.Run("Where_OnEmptyGroupedDataFrame", func(t *testing.T) { + fields := []arrow.Field{{Name: "key", Type: arrow.BinaryTypes.String}} + keyDataEmpty := getTestStringArray(mem, []string{}, nil); defer keyDataEmpty.Release() + emptyBaseDf, arrs := getBaseTestDfForGrouping(t, mem, "empty_df_where", fields, keyDataEmpty) + defer emptyBaseDf.(df.Releaser).Release() + for _, arr := range arrs { defer arr.Release() } + groupedEmpty := emptyBaseDf.GroupBy("key") + defer groupedEmpty.(arrowimpl.Releaser).Release() + + filtered := groupedEmpty.Where(func(kr df.Row, gdf df.DataFrame) bool { + if gdf != nil { gdf.(df.Releaser).Release() } + return true + }) + defer filtered.(arrowimpl.Releaser).Release() + assert.Equal(t, int64(0), filtered.Len()) + }) +} + +func TestGroupedDataFrame_Map(t *testing.T) { + mem := memory.NewGoAllocator() + originalDf, groupedDf := getTestGroupedDataFrameForAgg(t, mem) + defer originalDf.(df.Releaser).Release() + defer groupedDf.(arrowimpl.Releaser).Release() + + t.Run("Map_SelectFirstRowOfEachGroup", func(t *testing.T) { + // The map function will take each group (as a DataFrame) and return a new DataFrame + // containing only the first row of that group. + // The resulting GroupedDataFrame should then effectively contain these first rows, + // still grouped by the original keys. + mapFunc := func(keyRow df.Row, groupSubDf df.DataFrame) df.DataFrame { + if groupSubDf.Len() == 0 { + // Return an empty DF with the same schema if the group is empty + return arrowimpl.NewArrowDataFrame(groupSubDf.Name()+"_map_empty", nil, groupSubDf.Schema().(*arrowimpl.ArrowDataFrameSchema)) + } + // Select the first row. Limit(0,1) + firstRowDf := groupSubDf.Limit(0, 1) + // Important: The returned DataFrame from mapFunc must be an *arrowDataFrame + // and it must be explicitly managed (released) if not returned to something that takes ownership. + // In this case, the caller of Map (the test) will own the final result. + // The intermediate DFs created here inside mapFunc and returned to the Map method + // will be handled by the Map method's implementation (e.g. it concatenates records). + return firstRowDf + } + + mappedGroupedDf := groupedDf.Map(mapFunc) + defer mappedGroupedDf.(arrowimpl.Releaser).Release() + + assert.Equal(t, groupedDf.Len(), mappedGroupedDf.Len(), "Number of groups should remain the same after Map") + + // Verify content: each group in mappedGroupedDf should contain exactly one row, + // which is the first row of the corresponding group in the original groupedDf. + mappedGroupedDf.ForEach(func(keyRow df.Row, mappedSubGroupDf df.DataFrame) { + defer mappedSubGroupDf.(df.Releaser).Release() + assert.Equal(t, int64(1), mappedSubGroupDf.Len(), fmt.Sprintf("Group for key %v should have 1 row after map", keyRow)) + + originalSubGroupDf := groupedDf.Get(keyRow) // Get original group + defer originalSubGroupDf.(df.Releaser).Release() + + if originalSubGroupDf.Len() > 0 { + expectedFirstRow := originalSubGroupDf.GetRow(0) + actualFirstRowInMapped := mappedSubGroupDf.GetRow(0) + + // Compare row contents (example for first column) + // This requires a deep comparison of row values. + // For simplicity, we'll just check one value as a proxy. + // A full test would iterate all columns. + assert.Equal(t, expectedFirstRow.Get(2).Get(), actualFirstRowInMapped.Get(2).Get(), // Compare val_sum_int + fmt.Sprintf("Content mismatch for key %v", keyRow)) + } + }) + }) + + t.Run("Map_PanicIfFuncReturnsNil", func(t *testing.T) { + mapFuncNil := func(keyRow df.Row, groupSubDf df.DataFrame) df.DataFrame { + // Release groupSubDf as it's an input to this lambda and won't be used if we return nil + if r, ok := groupSubDf.(df.Releaser); ok { r.Release() } + return nil + } + assert.Panics(t, func() { + res := groupedDf.Map(mapFuncNil) + if res != nil { res.(arrowimpl.Releaser).Release() } + }) + }) + + t.Run("Map_PanicIfFuncReturnsNonArrowDf", func(t *testing.T) { + // Mock a non-arrow DataFrame + type dummyDataFrame struct { df.DataFrame } // Minimal struct to satisfy interface + // Add methods to dummyDataFrame to satisfy df.DataFrame if needed, or ensure type assertion is the primary check. + // For this test, the type assertion in arrowGroupedDataFrame.Map is key. + + mapFuncNonArrow := func(keyRow df.Row, groupSubDf df.DataFrame) df.DataFrame { + if r, ok := groupSubDf.(df.Releaser); ok { r.Release() } // Release the input group + return &dummyDataFrame{} // Return a non-arrow DataFrame + } + assert.Panics(t, func() { + res := groupedDf.Map(mapFuncNonArrow) + if res != nil { res.(arrowimpl.Releaser).Release() } + }) + }) +} + +func TestGroupedDataFrame_Agg(t *testing.T) { + mem := memory.NewGoAllocator() + originalDf, groupedDf := getTestGroupedDataFrameForAgg(t, mem) + defer originalDf.(df.Releaser).Release() + defer groupedDf.(arrowimpl.Releaser).Release() + + // Expected groups and their characteristics for manual verification: + // G1: {"groupA", 10} -> 3 rows. val_sum_int: {100, 200, nil} -> sum 300, mean 150. val_mean_float: {10.1, 20.2, 30.3} -> sum 60.6, mean 20.2. val_count_str_nullable: {"apple", "banana", nil} -> count 2 + // G2: {"groupB", 20} -> 2 rows. val_sum_int: {50, 60} -> sum 110, mean 55. val_mean_float: {5.5, 6.6} -> sum 12.1, mean 6.05. val_count_str_nullable: {"cat", "dog"} -> count 2 + // G3: {"groupA", 30} -> 1 row. val_sum_int: {70} -> sum 70, mean 70. val_mean_float: {7.7} -> sum 7.7, mean 7.7. val_count_str_nullable: {"eel"} -> count 1 + // G4: {nil, 10} -> 1 row. val_sum_int: {80} -> sum 80, mean 80. val_mean_float: {8.8} -> sum 8.8, mean 8.8. val_count_str_nullable: {"frog"} -> count 1 + // G5: {"groupC", 40} -> 1 row. val_sum_int: {90} -> sum 90, mean 90. val_mean_float: {9.9} -> sum 9.9, mean 9.9. val_count_str_nullable: {nil} -> count 0 + + t.Run("Agg_SingleCountSpecificColumn", func(t *testing.T) { + aggConfigs := []df.AggregationConfig{ + {Func: "count", InputCol: "val_count_str_nullable", OutputColName: "count_val_str"}, + } + aggDf := groupedDf.Agg(aggConfigs...) + defer aggDf.(df.Releaser).Release() + + assert.Equal(t, groupedDf.Len(), aggDf.Len(), "Number of rows in aggregated DF should equal number of groups") + assert.Equal(t, 3, aggDf.Schema().Len(), "Schema should have key cols + 1 agg col") // key1, key2, count_val_str + assert.Equal(t, "count_val_str", aggDf.Schema().Get(2).Name) + assert.Equal(t, df.IntegerFormat, aggDf.Schema().Get(2).Format) // Count is Int + + // Verify data - requires finding the correct row as order is not guaranteed. + // For simplicity, we'll check a known group, e.g., {"groupA", 10} + // This is brittle if GetRow order changes. A map-based check would be better for full validation. + // Let's find the row for {"groupA", 10} + var foundG1 bool + for i := int64(0); i < aggDf.Len(); i++ { + row := aggDf.GetRow(i) + if !row.GetByName("key1_str").IsNil() && row.GetByName("key1_str").GetAsString() == "groupA" && row.GetByName("key2_int").GetAsInt() == 10 { + assert.Equal(t, int64(2), row.GetByName("count_val_str").GetAsInt(), "Count for groupA,10") + foundG1 = true + break + } + } + assert.True(t, foundG1, "Group G1 (groupA, 10) not found in agg results") + }) + + t.Run("Agg_SingleCountAll", func(t *testing.T) { + aggConfigs := []df.AggregationConfig{ + {Func: "count", OutputColName: "group_row_count"}, // No InputCol + } + aggDf := groupedDf.Agg(aggConfigs...) + defer aggDf.(df.Releaser).Release() + assert.Equal(t, 3, aggDf.Schema().Len()) + assert.Equal(t, "group_row_count", aggDf.Schema().Get(2).Name) + + // Group {"groupA", 10} had 3 rows + var foundG1 bool + for i := int64(0); i < aggDf.Len(); i++ { + row := aggDf.GetRow(i) + if !row.GetByName("key1_str").IsNil() && row.GetByName("key1_str").GetAsString() == "groupA" && row.GetByName("key2_int").GetAsInt() == 10 { + assert.Equal(t, int64(3), row.GetByName("group_row_count").GetAsInt(), "Row count for groupA,10") + foundG1 = true; break + } + } + assert.True(t, foundG1) + }) + + t.Run("Agg_SingleSumInt", func(t *testing.T) { + aggConfigs := []df.AggregationConfig{{Func: "sum", InputCol: "val_sum_int", OutputColName: "sum_val_int"}} + aggDf := groupedDf.Agg(aggConfigs...) + defer aggDf.(df.Releaser).Release() + // val_sum_int for {"groupA", 10} is {100, 200, nil}, sum should be 300 + var foundG1 bool + for i := int64(0); i < aggDf.Len(); i++ { + row := aggDf.GetRow(i) + if !row.GetByName("key1_str").IsNil() && row.GetByName("key1_str").GetAsString() == "groupA" && row.GetByName("key2_int").GetAsInt() == 10 { + // Arrow's sum on int with nulls might produce int or float. Assuming int if all inputs are int. + // The actual type depends on compute kernel. Let's check the value. + // If original column was int64, sum is often int64. + assert.Equal(t, df.IntegerFormat, row.GetByName("sum_val_int").Schema().Format) + assert.Equal(t, int64(300), row.GetByName("sum_val_int").GetAsInt()) + foundG1 = true; break + } + } + assert.True(t, foundG1) + }) + + t.Run("Agg_SingleMeanFloat", func(t *testing.T) { + aggConfigs := []df.AggregationConfig{{Func: "mean", InputCol: "val_mean_float", OutputColName: "mean_val_float"}} + aggDf := groupedDf.Agg(aggConfigs...) + defer aggDf.(df.Releaser).Release() + // val_mean_float for {"groupA", 10} is {10.1, 20.2, 30.3}, mean should be 20.2 + var foundG1 bool + for i := int64(0); i < aggDf.Len(); i++ { + row := aggDf.GetRow(i) + if !row.GetByName("key1_str").IsNil() && row.GetByName("key1_str").GetAsString() == "groupA" && row.GetByName("key2_int").GetAsInt() == 10 { + assert.Equal(t, df.DoubleFormat, row.GetByName("mean_val_float").Schema().Format) // Mean is usually float + assert.InDelta(t, 20.2, row.GetByName("mean_val_float").GetAsDouble(), 0.00001) + foundG1 = true; break + } + } + assert.True(t, foundG1) + }) + + t.Run("Agg_MultipleAggregations", func(t *testing.T) { + aggConfigs := []df.AggregationConfig{ + {Func: "sum", InputCol: "val_sum_int", OutputColName: "total_sum_int"}, + {Func: "mean", InputCol: "val_mean_float", OutputColName: "avg_mean_float"}, + {Func: "count", InputCol: "val_count_str_nullable", OutputColName: "non_null_count_str"}, + {Func: "min", InputCol: "val_sum_int", OutputColName: "min_sum_int"}, + {Func: "max", InputCol: "val_mean_float", OutputColName: "max_mean_float"}, + } + aggDf := groupedDf.Agg(aggConfigs...) + defer aggDf.(df.Releaser).Release() + + assert.Equal(t, 2 + len(aggConfigs), aggDf.Schema().Len(), "Schema length for multiple aggs") + + // Spot check for group {"groupA", 10} + // SumInt: 300, MeanFloat: 20.2, CountStr: 2, MinSumInt: 100, MaxMeanFloat: 30.3 + var foundG1 bool + for i := int64(0); i < aggDf.Len(); i++ { + row := aggDf.GetRow(i) + if !row.GetByName("key1_str").IsNil() && row.GetByName("key1_str").GetAsString() == "groupA" && row.GetByName("key2_int").GetAsInt() == 10 { + assert.Equal(t, int64(300), row.GetByName("total_sum_int").GetAsInt()) + assert.InDelta(t, 20.2, row.GetByName("avg_mean_float").GetAsDouble(), 0.00001) + assert.Equal(t, int64(2), row.GetByName("non_null_count_str").GetAsInt()) + assert.Equal(t, int64(100), row.GetByName("min_sum_int").GetAsInt()) + assert.InDelta(t, 30.3, row.GetByName("max_mean_float").GetAsDouble(), 0.00001) + foundG1 = true; break + } + } + assert.True(t, foundG1) + }) + + t.Run("Agg_AllNullsColumn", func(t *testing.T){ + aggConfigs := []df.AggregationConfig{ + {Func: "sum", InputCol: "val_all_nulls_float", OutputColName: "sum_all_null"}, + {Func: "count", InputCol: "val_all_nulls_float", OutputColName: "count_all_null"}, + {Func: "mean", InputCol: "val_all_nulls_float", OutputColName: "mean_all_null"}, + } + aggDf := groupedDf.Agg(aggConfigs...) + defer aggDf.(df.Releaser).Release() + + for i := int64(0); i < aggDf.Len(); i++ { + row := aggDf.GetRow(i) + // Sum of all nulls is often null or 0 depending on kernel. Arrow sum kernel returns null if all inputs are null. + assert.True(t, row.GetByName("sum_all_null").IsNil(), "Sum of all nulls should be null") + assert.Equal(t, int64(0), row.GetByName("count_all_null").GetAsInt(), "Count of all nulls should be 0") + assert.True(t, row.GetByName("mean_all_null").IsNil(), "Mean of all nulls should be null") + } + }) + + t.Run("Agg_NoConfigsReturnsDistinctKeys", func(t *testing.T){ + aggDf := groupedDf.Agg() // No aggregation configs + defer aggDf.(df.Releaser).Release() + + assert.Equal(t, groupedDf.Len(), aggDf.Len(), "Agg with no configs should return same number of rows as groups") + assert.Equal(t, len(groupedDf.GetGroupColumns()), aggDf.Schema().Len(), "Schema should only contain group key columns") + + // Verify keys are the same + keysFromGrouped := groupedDf.GetKeys() + keysFromAgg := make([][]interface{}, aggDf.Len()) + keySchema := keysFromGrouped[0].Schema() // Assume at least one key + + for i:=int64(0); i 0 { + keyRowSchema := keys[0].Schema() + assert.Equal(t, len(expectedKeySchemaFields), keyRowSchema.Len(), "Key row schema length mismatch") + assert.Equal(t, "key1_str", keyRowSchema.Get(0).Name) + assert.Equal(t, df.StringFormat.Name(), keyRowSchema.Get(0).Format.Name()) // Check format name + assert.Equal(t, "key2_int", keyRowSchema.Get(1).Name) + assert.Equal(t, df.IntegerFormat.Name(), keyRowSchema.Get(1).Format.Name()) + } + + // Convert keys to a slice of slices for easier comparison after sorting + // Note: order of keys from GetKeys is not guaranteed. + actualKeyData := make([][]interface{}, len(keys)) + for i, keyRow := range keys { + actualKeyData[i] = []interface{}{keyRow.Get(0).Get(), keyRow.Get(1).Get()} + // Handle nil for string key for consistent sorting/comparison + if keyRow.Get(0).IsNil() { + actualKeyData[i][0] = nilPlaceholder // Use placeholder for sorting if actual nil is problematic + } + } + sortSliceOfInterfaceSlices(actualKeyData) + + expectedKeyData := [][]interface{}{ + {"groupA", int64(10)}, + {"groupA", int64(30)}, + {"groupB", int64(20)}, + {"groupC", int64(40)}, + {nilPlaceholder, int64(10)}, // Group with NULL key1_str + } + sortSliceOfInterfaceSlices(expectedKeyData) + assert.Equal(t, expectedKeyData, actualKeyData, "Key data mismatch") + }) + + t.Run("GetKeysOnEmptyGroupedDataFrame", func(t *testing.T) { + fields := []arrow.Field{{Name: "key", Type: arrow.BinaryTypes.String}} + keyDataEmpty := getTestStringArray(mem, []string{}, nil); defer keyDataEmpty.Release() + + emptyDf, arrs := getBaseTestDfForGrouping(t, mem, "empty_df_keys", fields, keyDataEmpty) + defer emptyDf.(df.Releaser).Release() + for _, arr := range arrs { defer arr.Release() } + + groupedEmpty := emptyDf.GroupBy("key") + defer groupedEmpty.(arrowimpl.Releaser).Release() + + keys := groupedEmpty.GetKeys() + assert.Empty(t, keys, "GetKeys on an empty grouped DataFrame should return empty slice.") + }) +} diff --git a/df/arrow/series.go b/df/arrow/series.go new file mode 100644 index 0000000..103d5b4 --- /dev/null +++ b/df/arrow/series.go @@ -0,0 +1,1017 @@ +//go:build arrow + +package arrow + +import ( + "context" + "fmt" + // "reflect" // Not used in the original, check if needed by my changes + // "time" // Not used in the original, check if needed by my changes + + "git.querycap.com/practice/df" // MODIFIED: Import path + "github.com/apache/arrow/go/v14/arrow" + "github.com/apache/arrow/go/v14/arrow/array" + "github.com/apache/arrow/go/v14/arrow/builder" // Using generic builder + "github.com/apache/arrow/go/v14/arrow/compute" + "github.com/apache/arrow/go/v14/arrow/memory" + "github.com/apache/arrow/go/v14/arrow/scalar" +) + +type arrowSeries struct { + schema df.SeriesSchema + arr arrow.Array + mem memory.Allocator +} + +// REMOVED local dfFormatToArrowType, will use the one from types.go +// REMOVED local appendScalarToBuilder, will use the one from types.go +// The dfValueToArrowScalar used by AsFormat, WhenNil, When will come from types.go + +func NewArrowSeries(arr arrow.Array, schema df.SeriesSchema) df.Series { return NewArrowSeriesWithAllocator(arr, schema, memory.DefaultAllocator) } +func NewArrowSeriesWithAllocator(rawArr arrow.Array, schema df.SeriesSchema, mem memory.Allocator) df.Series { + if rawArr == nil { panic("NewArrowSeriesWithAllocator: arrow.Array cannot be nil") }; + if mem == nil { panic("NewArrowSeriesWithAllocator: memory.Allocator cannot be nil") } + + // If schema.Format is UnknownFormat (or nil, if Format is an interface and nil is possible for "unknown"), + // and inference from array also yields UnknownFormat, then it's an issue, unless it's a NullType array. + if (schema.Format == df.UnknownFormat || schema.Format == nil) && + arrowToDfFormat(rawArr.DataType()) == df.UnknownFormat && + rawArr.DataType().ID() != arrow.NULL { + panic(fmt.Sprintf("NewArrowSeriesWithAllocator: df.SeriesSchema.Format is %v and cannot be inferred from array type %s", schema.Format, rawArr.DataType().Name())) + } + + // 1. Determine effective df.Format + effectiveFormat := schema.Format + if effectiveFormat == df.UnknownFormat || effectiveFormat == nil { + inferredFormat := arrowToDfFormat(rawArr.DataType()) + // If still unknown after inference (and not a Null array type), then panic. + if inferredFormat == df.UnknownFormat && rawArr.DataType().ID() != arrow.NULL { + panic(fmt.Sprintf("NewArrowSeriesWithAllocator: cannot infer df.Format from arrow array type %s and input schema.Format was %v", rawArr.DataType().Name(), schema.Format)) + } + effectiveFormat = inferredFormat + } + // After inference, effectiveFormat must be valid (not Unknown or nil if Format is an interface) + // unless it's a NullType array which can pair with UnknownFormat. + if (effectiveFormat == df.UnknownFormat || effectiveFormat == nil) && rawArr.DataType().ID() != arrow.NULL { + panic(fmt.Sprintf("NewArrowSeriesWithAllocator: effective df.Format is still %v after inference for non-Null array type %s", effectiveFormat, rawArr.DataType().Name())) + } + + // 2. Validate df.Format compatibility with rawArr.DataType() + // expectedArrowType is what the effectiveFormat maps to in Arrow terms. + expectedArrowType, err := dfFormatToArrowType(effectiveFormat) + if err != nil { + // This panic should ideally not be reached if effectiveFormat is not UnknownFormat + // and dfFormatToArrowType covers all known df.Formats. + panic(fmt.Sprintf("NewArrowSeriesWithAllocator: cannot map effectiveFormat %v to arrow.DataType: %v", effectiveFormat, err)) + } + + // If the actual array type doesn't match what's expected by effectiveFormat, + // check for known compatible groups (e.g. any Arrow int type vs df.IntFormat). + if rawArr.DataType().ID() != arrow.NULL && expectedArrowType.ID() != rawArr.DataType().ID() { + actualArrFormat := arrowToDfFormat(rawArr.DataType()) // What df.Format the actual array maps to + compatible := false + if isIntegerFormat(actualArrFormat) && isIntegerFormat(effectiveFormat) { + compatible = true + } else if isFloatFormat(actualArrFormat) && isFloatFormat(effectiveFormat) { + compatible = true + } else if isStringFormat(actualArrFormat) && isStringFormat(effectiveFormat) { + compatible = true + } else if isBoolFormat(actualArrFormat) && isBoolFormat(effectiveFormat) { + compatible = true + } else if isTimeFormat(actualArrFormat) && isTimeFormat(effectiveFormat) { + compatible = true + } else if isDateFormat(actualArrFormat) && isDateFormat(effectiveFormat) { + compatible = true + } + // Add other compatibility rules if necessary + + if !compatible && actualArrFormat != effectiveFormat { + panic(fmt.Sprintf("NewArrowSeriesWithAllocator: effectiveFormat %v (expects Arrow %s) is not compatible with actual array type %s (which maps to df.Format %v)", + effectiveFormat, expectedArrowType.Name(), rawArr.DataType().Name(), actualArrFormat)) + } + } + + // 3. Determine effective Nullability + effectiveNullable := schema.Nullable // Start with user's preference from input schema + if rawArr.NullN() > 0 { + // Data has nulls. If schema was explicitly set by user to NOT nullable, this is a contradiction. + // Check if schema.Format was originally set (not inferred) to make this decision. + if (schema.Format != nil && schema.Format != df.UnknownFormat) && !schema.Nullable { + panic(fmt.Sprintf("NewArrowSeriesWithAllocator: schema for series '%s' (Format: %v) is marked non-nullable, but data contains %d nulls", + schema.Name, effectiveFormat, rawArr.NullN())) + } + effectiveNullable = true // Data has nulls, so series must be nullable. + } else { + // Data has no nulls. + // If schema.Format was NOT specified by user (i.e., it was inferred), then nullability also gets inferred as false (non-nullable). + // Otherwise (schema.Format was specified by user), respect schema.Nullable specified by the user. + if schema.Format == nil || schema.Format == df.UnknownFormat { + effectiveNullable = false + } + // if schema.Format was specified by user, effectiveNullable is already schema.Nullable from above, which is correct. + } + + finalSchema := df.SeriesSchema{ + Name: schema.Name, + Format: effectiveFormat, + Nullable: effectiveNullable, + Metadata: schema.Metadata, // Preserve metadata + } + + rawArr.Retain() + // Use finalSchema here + return &arrowSeries{schema: finalSchema, arr: rawArr, mem: mem} +} + +func (as *arrowSeries) Schema() df.SeriesSchema { return as.schema } +func (as *arrowSeries) Len() int { if as.arr == nil { return 0 }; return as.arr.Len() } // MODIFIED: int64 to int + +func (as *arrowSeries) Get(index int) df.Value { // MODIFIED: int64 to int + if as.arr == nil || index < 0 || index >= as.arr.Len() { + // Return a typed null value consistent with Series format + arrowDt, err := dfFormatToArrowType(as.schema.Format) + if err != nil { + // This case is tricky: Get interface doesn't return error. + // Panic if we can't even determine the type for a null value. + panic(fmt.Errorf("Get: cannot determine arrow type for schema format %v for out-of-bounds access: %w", as.schema.Format, err)) + } + nullScalar := scalar.NewNullScalar(arrowDt) + return NewArrowValue(nullScalar, as.schema.Format) + } + return NewArrowValue(scalar.MakeScalar(as.arr, index), as.schema.Format) +} + +// REMOVED dfFormatToArrowTypeUnsafe helper, direct error handling or panic inline + +func (as *arrowSeries) ForEach(f func(df.Value)) { if as.arr == nil { return }; for i := 0; i < as.Len(); i++ { f(as.Get(i)) } } + +func (as *arrowSeries) Limit(offset int, size int) df.Series { + if as.arr == nil { panic("cannot limit a nil series") }; currentLen := as.arr.Len(); if offset < 0 { offset = 0 } + if offset >= currentLen { + // Return empty series of the same type + b := array.NewBuilder(as.mem, as.arr.DataType()); defer b.Release() // MODIFIED: builder.NewBuilder to array.NewBuilder + newArr := b.NewArray() + // newArr is already retained by NewArray, but NewArrowSeriesWithAllocator will retain again. + // This is fine as long as releases match. + return NewArrowSeriesWithAllocator(newArr, as.schema, as.mem) + } + if offset+size > currentLen { size = currentLen - offset }; if size < 0 { size = 0 } + newSlice := array.NewSlice(as.arr, int64(offset), int64(offset+size)) + // newSlice is retained. NewArrowSeriesWithAllocator will retain again. + return NewArrowSeriesWithAllocator(newSlice, as.schema, as.mem) +} + +func (as *arrowSeries) Where(f func(df.Value) bool) df.Series { + if as.arr == nil { panic("cannot filter a nil series") }; + b := array.NewBuilder(as.mem, as.arr.DataType()); defer b.Release() // MODIFIED: builder.NewBuilder to array.NewBuilder + for i := 0; i < as.Len(); i++ { // MODIFIED: int64 to int + val := as.Get(i) + if f(val) { + arrowVal, ok := val.(*arrowValue) + if !ok && !val.IsNil() { // If val is not nil, it must be arrowValue + panic(fmt.Errorf("Where: df.Value is not an arrowValue and not nil, type: %T", val)) + } + + var scalarToAppend scalar.Scalar + if val.IsNil() { // If df.Value itself is nil + scalarToAppend = scalar.NewNullScalar(b.Type()) + } else { // val is a valid arrowValue (which might wrap a nil scalar) + scalarToAppend = arrowVal.s + } + + err := appendScalarToBuilder(b, scalarToAppend, b.Type()) // Using types.go helper + if err != nil { + panic(fmt.Errorf("Where: error appending scalar to builder: %w. Scalar: %v, BuilderType: %v", err, scalarToAppend, b.Type())) + } + } + } + newArr := b.NewArray() + return NewArrowSeriesWithAllocator(newArr, as.schema, as.mem) +} + +func (as *arrowSeries) Sort(order df.SortOrder) df.Series { + if as.arr == nil || as.arr.Len() == 0 { return as.Copy() }; + ctx := compute.WithAllocator(context.Background(), as.mem) // compute.DefaultContext might be enough if no specific allocator needed for compute ops + + arrowSortOrder := arrow.Ascending + if order == df.Descending { // Assuming df.Descending is the correct enum value + arrowSortOrder = arrow.Descending + } + + // compute.SortIndices takes an Datum input + arrDatum := arrow.NewArrayDatum(as.arr) + // arrDatum does not need Release() if it's just a wrapper and doesn't retain as.arr + // However, Arrow examples often show it. Let's assume it's safer with Release if created. + // According to docs, NewArrayDatum does not acquire ownership (no Retain). So no Release needed. + + indicesDatum, err := compute.SortIndices(ctx, arrDatum, compute.SortOptions{Order: arrowSortOrder, NullPlacement: arrow.NullsFirst}) + if err != nil { panic(fmt.Errorf("SortIndices failed: %w", err)) }; + defer indicesDatum.Release() // indicesDatum is a new Datum, needs release. + + indicesArr, ok := indicesDatum.Value().(arrow.Array) // Use Value() then type assert + if !ok { panic("SortIndices did not return a valid ArrayDatum containing an Array") } + // indicesArr is owned by indicesDatum. + + // Create a new Datum for indicesArr for the Take operation. + // This might be redundant if Take can accept arrow.Array directly, but API expects Datum. + indicesArrDatum := arrow.NewArrayDatum(indicesArr) + // Not releasing indicesArrDatum as it's just a wrapper for indicesArr which is owned by indicesDatum. + + sortedArrDatum, err := compute.Take(ctx, compute.TakeOptions{}, arrDatum, indicesArrDatum) + if err != nil { panic(fmt.Errorf("Take failed: %w", err)) }; + defer sortedArrDatum.Release() // sortedArrDatum is new, needs release. + + sortedArr, ok := sortedArrDatum.Value().(arrow.Array) + if !ok { panic("Take did not return a valid ArrayDatum containing an Array") } + // sortedArr is owned by sortedArrDatum. NewArrowSeriesWithAllocator will Retain. + return NewArrowSeriesWithAllocator(sortedArr, as.schema, as.mem) +} + +func (as *arrowSeries) Map(outputFormat df.Format, f func(df.Value) df.Value) df.Series { + if as.arr == nil { panic("map on nil series") }; + outputArrowType, err := dfFormatToArrowType(outputFormat) + if err != nil { + panic(fmt.Errorf("Map: could not get arrow type for output format %v: %w", outputFormat, err)) + } + + b := array.NewBuilder(as.mem, outputArrowType); defer b.Release() + for i := 0; i < as.Len(); i++ { + originalVal := as.Get(i); + mappedVal := f(originalVal) // mappedVal is df.Value + + scalarToAppend, errConv := dfValueToArrowScalar(mappedVal, outputArrowType) + if errConv != nil { + panic(fmt.Errorf("Map: failed to convert mapped df.Value (from input %v) to Arrow scalar for type %s: %w", originalVal, outputArrowType.Name(), errConv)) + } + // dfValueToArrowScalar from types.go does not retain the scalar it returns, so no release needed for scalarToAppend here. + + errAppend := appendScalarToBuilder(b, scalarToAppend, outputArrowType) + if errAppend != nil { + panic(fmt.Errorf("Map: error appending scalar to builder: %w. Scalar: %v", errAppend, scalarToAppend)) + } + } + newArr := b.NewArray() + finalNullable := true // Safest default: operations might introduce nulls. + if newArr.NullN() == 0 { + // If no nulls were produced, can we infer it's non-nullable? + // This is only true if the mapping function `f` guarantees non-null output + // AND the outputFormat itself isn't something like a "nullable string" type. + // This level of inference is hard. A simpler rule: + // If the original series was non-nullable AND f guarantees non-null, then non-nullable. + // For now, if newArr.NullN() == 0, we *could* set it to false, but it's an assumption. + // Let's assume `f` can produce nils, so `true` is safer unless `newArr.NullN() == 0`. + // If `f` *cannot* produce nils (e.g. `v.GetAsInt() * 2`), and input has no nils for type errors, + // then nullability might be preserved or become false. + // Sticking to `newArr.NullN() == 0` is a data-driven way. + finalNullable = false + } + newSchema := df.SeriesSchema{Name: as.schema.Name, Format: outputFormat, Nullable: finalNullable, Metadata: as.schema.Metadata} + return NewArrowSeriesWithAllocator(newArr, newSchema, as.mem) +} + +func (as *arrowSeries) FlatMap(outputFormat df.Format, f func(df.Value) []df.Value) df.Series { + if as.arr == nil { panic("flatMap on nil series") }; + outputArrowType, err := dfFormatToArrowType(outputFormat) + if err != nil { + panic(fmt.Errorf("FlatMap: could not get arrow type for output format %v: %w", outputFormat, err)) + } + + b := array.NewBuilder(as.mem, outputArrowType); defer b.Release() + for i := 0; i < as.Len(); i++ { + originalVal := as.Get(i) + for _, mappedVal := range f(originalVal) { // mappedVal is df.Value + scalarToAppend, errConv := dfValueToArrowScalar(mappedVal, outputArrowType) + if errConv != nil { + panic(fmt.Errorf("FlatMap: failed to convert mapped df.Value (from input %v) to Arrow scalar for type %s: %w", originalVal, outputArrowType.Name(), errConv)) + } + // No release for scalarToAppend from dfValueToArrowScalar (from types.go) + + errAppend := appendScalarToBuilder(b, scalarToAppend, outputArrowType) + if errAppend != nil { + panic(fmt.Errorf("FlatMap: error appending scalar to builder: %w. Scalar: %v", errAppend, scalarToAppend)) + } + } + } + newArr := b.NewArray() + finalNullable := true // FlatMap can easily produce more or fewer items; nulls can appear. + if newArr.NullN() == 0 { + // Similar to Map, if f guarantees non-null, and output type isn't inherently nullable, + // then Nullable could be false. + finalNullable = false + } + newSchema := df.SeriesSchema{Name: as.schema.Name, Format: outputFormat, Nullable: finalNullable, Metadata: as.schema.Metadata} + return NewArrowSeriesWithAllocator(newArr, newSchema, as.mem) +} + +func (as *arrowSeries) Reduce(f func(df.Value, df.Value) df.Value, startValue df.Value) df.Value { + if startValue == nil { panic("Reduce startValue cannot be nil interface") }; + acc := startValue + if as.arr == nil || as.Len() == 0 { return acc } + for i := 0; i < as.Len(); i++ { acc = f(acc, as.Get(i)) } // MODIFIED: int64 to int + return acc +} + +func (as *arrowSeries) Distinct() df.Series { + if as.arr == nil || as.arr.Len() == 0 { return as.Copy() } + ctx := compute.WithAllocator(context.Background(), as.mem) // Or compute.DefaultContext() + arrDatum := arrow.NewArrayDatum(as.arr); defer arrDatum.Release() + uniqueDatum, err := compute.Unique(ctx, arrDatum) + if err != nil { panic(fmt.Sprintf("Unique failed: %v", err)) }; + defer uniqueDatum.Release() + + uniqueArr, ok := uniqueDatum.(*arrow.ArrayDatum).Value().(arrow.Array) + if !ok { panic("Unique did not return a valid ArrayDatum containing an Array") } + return NewArrowSeriesWithAllocator(uniqueArr, as.schema, as.mem) +} + +func (as *arrowSeries) Copy() df.Series { + if as.arr == nil { + if as.schema.Format != nil && as.mem != nil { + dt, err := dfFormatToArrowType(as.schema.Format) + if err != nil { + panic(fmt.Errorf("Copy: cannot determine arrow type for nil series' schema format %v: %w", as.schema.Format, err)) + } + bld := array.NewBuilder(as.mem, dt); defer bld.Release() + emptyArr := bld.NewArray() // Retained + return NewArrowSeriesWithAllocator(emptyArr, as.schema, as.mem) + } + panic("cannot copy nil series with no type/allocator info") + } + // NewSlice creates a new array struct that shares data buffers but has its own ref count. + // This is a shallow copy of data, but a new array instance. + newSlice := array.NewSlice(as.arr, 0, int64(as.arr.Len())) + return NewArrowSeriesWithAllocator(newSlice, as.schema, as.mem) +} + +func (as *arrowSeries) Release() { if as.arr != nil { as.arr.Release(); as.arr = nil } } + +// prepareOtherForSetOp (from my template) is better than direct type assertion and panic. +// I'll replace the direct assertions in Append, Intersection, etc., with this helper or similar logic. +func (as *arrowSeries) prepareOtherForSetOp(otherRaw df.Series, operationName string) (arrow.Array, error) { + if otherRaw == nil { + return nil, fmt.Errorf("%s: other series cannot be nil", operationName) + } + other, ok := otherRaw.(*arrowSeries) + if !ok { + return nil, fmt.Errorf("%s: expected *arrowSeries, got %T. Cross-implementation operations not yet supported.", operationName, otherRaw) + } + if other.arr == nil { + return nil, fmt.Errorf("%s: other series has nil internal array", operationName) + } + + if arrow.TypeEqual(as.arr.DataType(), other.arr.DataType()) { + other.arr.Retain() // Caller must release + return other.arr, nil + } + + // Try to cast 'other' to 'as' type if df.Formats are the same (suggesting conceptual compatibility) + if as.schema.Format == other.schema.Format { + ctx := compute.DefaultContext() // Or compute.WithAllocator(context.Background(), as.mem) + castedOtherArray, castErr := compute.Cast(ctx, other.arr, as.arr.DataType(), compute.DefaultCastOptions(false)) + if castErr != nil { + return nil, fmt.Errorf("%s: type mismatch (self=%s, other=%s) and cast failed: %w", + operationName, as.arr.DataType().Name(), other.arr.DataType().Name(), castErr) + } + // castedOtherArray is new and needs to be managed by caller (usually released after use) + return castedOtherArray, nil + } + + return nil, fmt.Errorf("%s: type mismatch (self=%s, other=%s) and df.Format mismatch (self=%v, other=%v)", + operationName, as.arr.DataType().Name(), other.arr.DataType().Name(), as.schema.Format, other.schema.Format) +} + + +func (as *arrowSeries) Append(otherSeriesRaw df.Series) df.Series { + if as.arr == nil { // If current series is nil (e.g. uninitialized) + if otherSeriesRaw == nil || otherSeriesRaw.Len() == 0 { // Appending nothing to nil series + // Return a valid empty series or panic, current Copy handles creating empty from schema + return as.Copy() + } + // If as.arr is nil but other is not, effectively this becomes otherSeriesRaw.Copy() + // This case should ideally be handled by ensuring as.arr is initialized (e.g. to empty array) + // For now, let's assume if as.arr is nil, it's an empty series of its schema type. + // This means it should behave like Len() == 0. + } + + if otherSeriesRaw == nil || otherSeriesRaw.Len() == 0 { + return as.Copy() // Appending empty series is a no-op + } + if as.Len() == 0 { // If current series is empty but other is not + // Check type compatibility before just copying otherSeriesRaw + // to ensure the result is of type as.schema.Format + _, ok := otherSeriesRaw.(*arrowSeries) + if !ok { panic(fmt.Sprintf("Append: expected *arrowSeries, got %T", otherSeriesRaw))} + + // If otherSeriesRaw's format matches as.schema.Format, then copy is fine. + // Otherwise, it might need conversion. + if otherSeriesRaw.Schema().Format != as.schema.Format { + // This implies a conversion is needed for the "empty" part of 'as' + // which is complex. Simplest is to require format match or panic. + // Or, treat as otherSeriesRaw.AsFormat(as.schema.Format). + panic(fmt.Sprintf("Append: format mismatch when appending to empty series. Self: %v, Other: %v", as.schema.Format, otherSeriesRaw.Schema().Format)) + } + return otherSeriesRaw.Copy() + } + + otherArr, err := as.prepareOtherForSetOp(otherSeriesRaw, "Append") + if err != nil { panic(err) } + defer otherArr.Release() + + concatenatedArr, err := array.Concatenate([]arrow.Array{as.arr, otherArr}, as.mem) + if err != nil { panic(fmt.Sprintf("Append: failed to concatenate arrays: %v", err)) } + // concatenatedArr is retained by Concatenate. NewArrowSeriesWithAllocator will retain again. + return NewArrowSeriesWithAllocator(concatenatedArr, as.schema, as.mem) +} + +func (as *arrowSeries) Union(otherSeries df.Series) df.Series { + if as.arr == nil { panic("Union on nil series") } + // Appending will handle type checks and potential casting if formats match. + appended := as.Append(otherSeries) // This returns a new series + // Make sure to release the intermediate 'appended' series' array + defer appended.Release() + + // Distinct will operate on the result of Append. + return appended.Distinct() +} + +func (as *arrowSeries) Intersection(otherSeriesRaw df.Series) df.Series { + if as.arr == nil { panic("Intersection on nil series") } + if otherSeriesRaw == nil || otherSeriesRaw.Len() == 0 || as.Len() == 0 { + dt, err := dfFormatToArrowType(as.schema.Format) + if err != nil { panic(fmt.Errorf("Intersection: cannot get arrow type for empty result: %w", err)) } + bld := array.NewBuilder(as.mem, dt); defer bld.Release() + emptyArr := bld.NewArray(); // Retained + return NewArrowSeriesWithAllocator(emptyArr, as.schema, as.mem) + } + + otherArr, err := as.prepareOtherForSetOp(otherSeriesRaw, "Intersection") + if err != nil { panic(fmt.Errorf("Intersection prepare error: %w", err)) } + defer otherArr.Release() // otherArr was retained by prepareOtherForSetOp or is new (casted) + + ctx := compute.WithAllocator(context.Background(), as.mem) + // NewArrayData does not retain input array, so no release needed for leftDatum/rightDatum wrappers + leftDatum := arrow.NewArrayDatum(as.arr) + rightDatum := arrow.NewArrayDatum(otherArr) + + resultSetDatum, err := compute.SetIntersection(ctx, leftDatum, rightDatum, compute.SetLookupOptions{NullMatchingBehavior: compute.MatchNulls}) + if err != nil { panic(fmt.Sprintf("Intersection: compute.SetIntersection failed: %w", err)) }; + defer resultSetDatum.Release() // This is a new datum, release it + + resultArr, ok := resultSetDatum.Value().(arrow.Array) + if !ok { panic("Intersection: compute.SetIntersection did not return ArrayDatum containing an Array") } + // resultArr is owned by resultSetDatum. NewArrowSeries retains it. + return NewArrowSeriesWithAllocator(resultArr, as.schema, as.mem) +} + +func (as *arrowSeries) Except(otherSeriesRaw df.Series) df.Series { + if as.arr == nil { panic("Except on nil series") } + if as.Len() == 0 { + dt, err := dfFormatToArrowType(as.schema.Format) + if err != nil { panic(fmt.Errorf("Except: cannot get arrow type for empty result: %w", err)) } + bld := array.NewBuilder(as.mem, dt); defer bld.Release() + emptyArr := bld.NewArray(); // Retained + return NewArrowSeriesWithAllocator(emptyArr, as.schema, as.mem) + } + if otherSeriesRaw == nil || otherSeriesRaw.Len() == 0 { return as.Copy() } + + otherArr, err := as.prepareOtherForSetOp(otherSeriesRaw, "Except") + if err != nil { panic(fmt.Errorf("Except prepare error: %w", err)) } + defer otherArr.Release() // otherArr was retained by prepareOtherForSetOp or is new (casted) + + ctx := compute.WithAllocator(context.Background(), as.mem) + leftDatum := arrow.NewArrayDatum(as.arr) + rightDatum := arrow.NewArrayDatum(otherArr) + + resultSetDatum, err := compute.SetDifference(ctx, leftDatum, rightDatum, compute.SetLookupOptions{NullMatchingBehavior: compute.MatchNulls}) + if err != nil { panic(fmt.Errorf("Except: compute.SetDifference failed: %w", err)) }; + defer resultSetDatum.Release() // This is a new datum, release it + + resultArr, ok := resultSetDatum.Value().(arrow.Array) + if !ok { panic("Except: compute.SetDifference did not return ArrayDatum containing an Array") } + // resultArr is owned by resultSetDatum. NewArrowSeries retains it. + return NewArrowSeriesWithAllocator(resultArr, as.schema, as.mem) +} + +func (as *arrowSeries) AsFormat(targetFormat df.Format) df.Series { + if as.arr == nil { panic("AsFormat called on nil series array") } + if targetFormat == nil { panic("AsFormat: targetFormat cannot be nil") } // Assuming df.Format is non-nil interface + + if as.schema.Format == targetFormat { // TODO: Ensure df.Format has proper Equals method if it's not a basic comparable type. + return as.Copy() + } + + targetArrowType, err := dfFormatToArrowType(targetFormat) + if err != nil { + panic(fmt.Errorf("AsFormat: cannot map target df.Format %v to Arrow type: %w", targetFormat, err)) + } + + if arrow.TypeEqual(as.arr.DataType(), targetArrowType) { + // Data type is already correct, just update schema.Format + // NewArrowSeriesWithAllocator will retain as.arr + newSchema := df.SeriesSchema{Name: as.schema.Name, Format: targetFormat, Nullable: as.arr.NullN() > 0} + return NewArrowSeriesWithAllocator(as.arr, newSchema, as.mem) + } + + ctx := compute.WithAllocator(context.Background(), as.mem) + castOptions := compute.DefaultCastOptions(false) // false = allow unsafe casts like float to int truncation. + // Set to true for strict (error on overflow/truncation). + castedArray, err := compute.Cast(ctx, as.arr, targetArrowType, castOptions) + if err != nil { + panic(fmt.Errorf("AsFormat: failed to cast array from %s (df.Format %v) to %s (df.Format %v): %w", + as.arr.DataType().Name(), as.schema.Format, targetArrowType.Name(), targetFormat, err)) + } + // castedArray is new and retained by compute.Cast. NewArrowSeriesWithAllocator will retain again. + + newSchema := df.SeriesSchema{Name: as.schema.Name, Format: targetFormat, Nullable: castedArray.NullN() > 0, Metadata: as.schema.Metadata} + return NewArrowSeriesWithAllocator(castedArray, newSchema, as.mem) +} + + +func (as *arrowSeries) WhenNil(fillValue df.Value) df.Series { + if as.arr == nil { panic("WhenNil called on nil series array") } + if fillValue == nil { panic("WhenNil: fillValue interface cannot be nil (can be a nil df.Value though)")} // df.Value can be nil if interface itself is nil + + if as.arr.NullN() == 0 { return as.Copy() } // No nulls to fill + + ctx := compute.WithAllocator(context.Background(), as.mem) + targetArrowType := as.arr.DataType() + + fillScalar, err := dfValueToArrowScalar(fillValue, targetArrowType) + if err != nil { + panic(fmt.Errorf("WhenNil: error converting fillValue (value: %v, format: %v) to Arrow scalar for type %s: %w", + fillValue.Get(), fillValue.Schema().Format, targetArrowType.Name(), err)) + } + + // If fillValue itself was a nil df.Value, then fillScalar will be a !IsValid scalar (a typed Null). + // compute.FillNull correctly handles filling with a typed Null scalar, effectively making it a no-op + // for those nulls being filled with another null. + // The original check `if fillValue.IsNil() && targetArrowType.ID() != arrow.NULL` and then + // `if !fillScalar.IsValid() { return as.Copy() }` was trying to optimize for this. + // compute.FillNull is idempotent if fillScalar is a typed null of the array's type. + // So, no need for an early exit if fillValue.IsNil(). + + fillScalarDatum := arrow.NewScalarDatum(fillScalar) + seriesDatum := arrow.NewArrayDatum(as.arr) + + resultDatum, err := compute.FillNull(ctx, seriesDatum, fillScalarDatum) + if err != nil { panic(fmt.Errorf("WhenNil: FillNull compute failed: %w", err)) } + defer resultDatum.Release() + + // The resultDatum contains the new array. + newArr := resultDatum.Value().(arrow.Array) + + currentSchema := as.schema + finalNullable := true + if !fillValue.IsNil() && newArr.NullN() == 0 { // If filled with non-null and result has no nulls + finalNullable = false + } else if fillValue.IsNil() && as.arr.NullN() > 0 { // If filled with null, nullability depends on original + finalNullable = true // Still nullable + } else if newArr.NullN() > 0 { // If there are still nulls for any other reason + finalNullable = true + } else { // No nulls in result, and fillValue was not nil (already covered), or original had no nulls + finalNullable = false + } + + newSchema := df.SeriesSchema{ + Name: currentSchema.Name, + Format: currentSchema.Format, + Nullable: finalNullable, + Metadata: currentSchema.Metadata, + } + return NewArrowSeriesWithAllocator(newArr, newSchema, as.mem) +} + +func (as *arrowSeries) When(replacementMap map[any]df.Value) df.Series { + if as.arr == nil { panic("When called on nil series array") } + if len(replacementMap) == 0 { return as.Copy() } + + colType := as.arr.DataType() + b := array.NewBuilder(as.mem, colType); defer b.Release() // MODIFIED: builder.NewBuilder to array.NewBuilder + + for r := 0; r < as.Len(); r++ { + currentDfVal := as.Get(r) // This is *arrowValue from our Get method + + // Determine the key for map lookup. If df.Value is nil, use nil as key. + // Otherwise, use its Go representation. + var goKeyForLookup any + if currentDfVal.IsNil() { + goKeyForLookup = nil // Standard way to represent nil in a map key if desired + } else { + // arrowValue.Get() returns any. This should be fine for map keys if types are simple. + goKeyForLookup = currentDfVal.Get() + } + + replacementDfVal, shouldReplace := replacementMap[goKeyForLookup] + + var scalarToAppend scalar.Scalar + var errConv error + + if shouldReplace { + scalarToAppend, errConv = dfValueToArrowScalar(replacementDfVal, colType) + if errConv != nil { + panic(fmt.Errorf("When: error converting replacement df.Value (for key %v, value %v) to Arrow scalar type %s: %w", + goKeyForLookup, replacementDfVal, colType.Name(), errConv)) + } + } else { + // No replacement, use original scalar. + // currentDfVal must be *arrowValue to get .s + arrowVal, ok := currentDfVal.(*arrowValue) + if !ok && !currentDfVal.IsNil() { // If not nil, it must be arrowValue + panic(fmt.Errorf("When: original df.Value is not *arrowValue and not nil, type: %T", currentDfVal)) + } + if currentDfVal.IsNil() { // If original was nil, create a nil scalar of target type + scalarToAppend = scalar.NewNullScalar(colType) + } else { + scalarToAppend = arrowVal.s + } + } + // dfValueToArrowScalar from types.go does not retain, so no release for scalarToAppend here. + + errAppend := appendScalarToBuilder(b, scalarToAppend, colType) + if errAppend != nil { + panic(fmt.Errorf("When: error appending scalar (for key %v, scalar %v) to builder: %w", + goKeyForLookup, scalarToAppend, errAppend)) + } + } + newArr := b.NewArray() + + currentSchema := as.schema + finalNullable := true // Default to true because replacements can introduce nulls. + if newArr.NullN() == 0 { + // If no nulls in the new array, it *could* be non-nullable. + // This is true if the original array had no nulls that were *not* replaced, + // AND all replacement values used were non-nil. + allReplacementsWereNonNull := true + for _, valToReplaceWith := range replacementMap { + if valToReplaceWith.IsNil() { + allReplacementsWereNonNull = false + break + } + } + if allReplacementsWereNonNull { // If all potential replacements are non-null + // We still need to consider if an original non-null value that wasn't in the map + // could have been preserved. This is guaranteed. + // So, if newArr.NullN() == 0 AND all replacementMap values are non-Null, + // it implies the resulting series is non-nullable. + finalNullable = false + } + } + + newSchema := df.SeriesSchema{ + Name: currentSchema.Name, + Format: currentSchema.Format, // Type is preserved + Nullable: finalNullable, + Metadata: currentSchema.Metadata, + } + return NewArrowSeriesWithAllocator(newArr, newSchema, as.mem) +} + +// --- Stubs for remaining methods --- +func (as *arrowSeries) Expr() df.Expr { + // This method should return an expression that represents this series. + // Typically, this would be a column name expression. + // Assuming df.NewColExpr (or a similar constructor) exists in the df package + // that creates an expression representing a column. + // The actual implementation depends on how df.Expr is defined and constructed. + // For this example, let's assume df.NewColExpr takes the column name. + if as.schema.Name == "" { + // If the series doesn't have a name, it's hard to represent it as a simple column expression. + // This might indicate an anonymous series, which could be problematic for Expr(). + // Depending on df.Expr capabilities, could return a special type of expression + // or panic if a name is essential for a column expression. + panic("Expr: cannot create a column expression for an unnamed series") + } + return df.NewColExpr(as.schema.Name) // Example: uses a hypothetical constructor +} + +func (as *arrowSeries) Select(e df.Expr) df.Series { + // Implementing a full expression evaluation engine for a single series is complex. + // It would involve evaluating the expression `e` where `as` is the context. + // For example, if `e` is `Col("this_series_name").Add(Literal(5))`, + // it would add 5 to each element of `as`. + // Many common operations are already covered by Map, AsFormat, or direct compute. + // For now, as per subtask, this will remain partially implemented. + // panic(fmt.Sprintf("Select on arrowSeries is partially implemented. Full expression (%s) evaluation TBD.", e.Name())) + + ctx := compute.WithAllocator(context.Background(), as.mem) + var currentArr arrow.Array + var currentSeriesSchema df.SeriesSchema + var seriesToRelease df.Series // To release intermediate series from parent expressions + + if e.Parent() != nil { + // Recursively evaluate the parent expression first + parentSeries := as.Select(e.Parent()) // This series is the input to the current operation + arrowParentSeries, ok := parentSeries.(*arrowSeries) + if !ok { + if parentSeries != nil { parentSeries.Release() } + panic(fmt.Sprintf("Select: parent expression did not return *arrowSeries, got %T", parentSeries)) + } + currentArr = arrowParentSeries.arr // Do not retain here, it's owned by arrowParentSeries + currentSeriesSchema = arrowParentSeries.schema + seriesToRelease = arrowParentSeries // Mark for release after use + } else { + // This expression operates directly on the current series 'as' + currentArr = as.arr + currentSeriesSchema = as.schema + } + + if currentArr == nil { + if seriesToRelease != nil { seriesToRelease.Release() } + panic(fmt.Sprintf("Select: input array for expression '%s' is nil", e.Name())) + } + + // Retain currentArr for operations, as it might be from `as.arr` or a temporary parent result. + // The final result array will be new or a slice, so this specific retain is for this scope. + currentArr.Retain(); defer currentArr.Release() + + + var resultArray arrow.Array + var resultSchema df.SeriesSchema + + switch e.OpType() { + case df.ExprTypeMap: + mapOp := e.MapOp() + if mapOp == nil { + if seriesToRelease != nil { seriesToRelease.Release() } + panic(fmt.Sprintf("Select: MapOp is nil for ExprTypeMap expression '%s'", e.Name())) + } + + // Assumptions about MapOp structure: + // - mapOp.Name() might give "OpConst_Add", "WhenNilConst" etc. + // - mapOp.Args() gives arguments, where literal is often Args()[0] + // This is a simplification; a real impl might need type assertions on mapOp + // or more detailed methods on the MapOp interface. + + opName := "" // This needs to be derived from how MapOp is structured in your df.Expr + // For example, if MapOp has a Name() method or if e.Name() directly gives the map op: + if exprWithName, ok := mapOp.(interface{ Name() string }); ok { // Hypothetical + opName = exprWithName.Name() + } else if exprWithName, ok := e.(interface{ Name() string }); ok { // Fallback to expr name itself + opName = exprWithName.Name() + } else { + // Try to infer from common types if not directly named + // This part is highly speculative based on common patterns in your expr package. + // It's better if mapOp itself has a clear way to identify the operation. + // Example: if type is *expr.SeriesWhenNilConstMapOp -> "WhenNilConst" + // For now, let's assume it's part of e.Name() or mapOp.Name() + // This logic will likely need refinement based on actual df.Expr structure. + opName = e.Name() // Fallback, might need adjustment. + } + + + if strings.HasPrefix(opName, "OpConst_") { // Assuming names like "OpConst_Add", "OpConst_Subtract" + if len(mapOp.Args()) == 0 || mapOp.Args()[0].Const() == nil { + if seriesToRelease != nil { seriesToRelease.Release() } + panic(fmt.Sprintf("Select: MapOp '%s' for expression '%s' requires a literal argument", opName, e.Name())) + } + literalValue := mapOp.Args()[0].Const() + + // For arithmetic, scalar type should ideally match array type or be promotable. + // Using currentArr.DataType() as the target for the scalar conversion. + rightScalarVal, err := dfValueToArrowScalar(literalValue, currentArr.DataType()) + if err != nil { + if seriesToRelease != nil { seriesToRelease.Release() } + panic(fmt.Sprintf("Select: error converting literal for MapOp '%s' expr '%s': %v", opName, e.Name(), err)) + } + // No release for rightScalarVal from dfValueToArrowScalar + + inputDatum := arrow.NewArrayDatum(currentArr) + rightScalarDatum := arrow.NewScalarDatum(rightScalarVal) + defer inputDatum.Release(); defer rightScalarDatum.Release() // Release datums + + var computeErr error + var outputDatum arrow.Datum + + switch opName { + case "OpConst_Add": // This name needs to match what your df.Expr generates + outputDatum, computeErr = compute.Add(ctx, inputDatum, rightScalarDatum, compute.ArithmeticOptions{NoSignedOverflow: false}) + case "OpConst_Subtract": + outputDatum, computeErr = compute.Subtract(ctx, inputDatum, rightScalarDatum, compute.ArithmeticOptions{NoSignedOverflow: false}) + case "OpConst_Multiply": + outputDatum, computeErr = compute.Multiply(ctx, inputDatum, rightScalarDatum, compute.ArithmeticOptions{NoSignedOverflow: false}) + case "OpConst_Divide": + outputDatum, computeErr = compute.Divide(ctx, inputDatum, rightScalarDatum, compute.ArithmeticOptions{NoSignedOverflow: false}) + default: + if seriesToRelease != nil { seriesToRelease.Release() } + panic(fmt.Sprintf("Select: unsupported MapOp arithmetic operation '%s' for expression '%s'", opName, e.Name())) + } + + if computeErr != nil { + if seriesToRelease != nil { seriesToRelease.Release() } + panic(fmt.Sprintf("Select: compute error for MapOp '%s' expr '%s': %v", opName, e.Name(), computeErr)) + } + defer outputDatum.Release() + resultArray = outputDatum.MakeArray() // Makes a new array, caller owns. + resultSchema = currentSeriesSchema // Arithmetic ops usually preserve type & name. Name might change via expr.Name(). + if e.Name() != "" { resultSchema.Name = e.Name() } + + + } else if opName == "WhenNilConst" { // Assuming this is the name for WhenNil operation + if len(mapOp.Args()) == 0 || mapOp.Args()[0].Const() == nil { + if seriesToRelease != nil { seriesToRelease.Release() } + panic(fmt.Sprintf("Select: MapOp WhenNilConst for expression '%s' requires a literal fill value", e.Name())) + } + fillValue := mapOp.Args()[0].Const() + fillScalarVal, err := dfValueToArrowScalar(fillValue, currentArr.DataType()) + if err != nil { + if seriesToRelease != nil { seriesToRelease.Release() } + panic(fmt.Sprintf("Select: error converting fillValue for WhenNilConst expr '%s': %v", e.Name(), err)) + } + + inputDatum := arrow.NewArrayDatum(currentArr) + fillScalarDatum := arrow.NewScalarDatum(fillScalarVal) + defer inputDatum.Release(); defer fillScalarDatum.Release() + + outputDatum, fillErr := compute.FillNull(ctx, inputDatum, fillScalarDatum) + if fillErr != nil { + if seriesToRelease != nil { seriesToRelease.Release() } + panic(fmt.Sprintf("Select: FillNull compute error for expr '%s': %v", e.Name(), fillErr)) + } + defer outputDatum.Release() + resultArray = outputDatum.MakeArray() + resultSchema = currentSeriesSchema + resultSchema.Nullable = resultArray.NullN() > 0 // Update nullability + if e.Name() != "" { resultSchema.Name = e.Name() } + + } else { + if seriesToRelease != nil { seriesToRelease.Release() } + panic(fmt.Sprintf("Select: unsupported MapOp operation '%s' for expression '%s'", opName, e.Name())) + } + + case df.ExprTypeFilter: + filterOp := e.FilterOp() + if filterOp == nil { + if seriesToRelease != nil { seriesToRelease.Release() } + panic(fmt.Sprintf("Select: FilterOp is nil for ExprTypeFilter expression '%s'", e.Name())) + } + + opName := "" // Similar to MapOp, need a way to get operation name like "OpFilter_EqConst" + if exprWithName, ok := filterOp.(interface{ Name() string }); ok { // Hypothetical + opName = exprWithName.Name() + } else if exprWithName, ok := e.(interface{ Name() string }); ok { + opName = exprWithName.Name() + } else { + opName = e.Name() // Fallback + } + + + if strings.HasPrefix(opName, "OpFilter_") { // e.g. "OpFilter_EqConst" + if len(filterOp.Args()) == 0 || filterOp.Args()[0].Const() == nil { + if seriesToRelease != nil { seriesToRelease.Release() } + panic(fmt.Sprintf("Select: FilterOp '%s' for expression '%s' requires a literal argument", opName, e.Name())) + } + literalValue := filterOp.Args()[0].Const() + rightScalarVal, err := dfValueToArrowScalar(literalValue, currentArr.DataType()) + if err != nil { + if seriesToRelease != nil { seriesToRelease.Release() } + panic(fmt.Sprintf("Select: error converting literal for FilterOp '%s' expr '%s': %v", opName, e.Name(), err)) + } + + inputDatum := arrow.NewArrayDatum(currentArr) + rightScalarDatum := arrow.NewScalarDatum(rightScalarVal) + defer inputDatum.Release(); defer rightScalarDatum.Release() + + var compareOpt compute.CompareOptions + switch opName { + case "OpFilter_EqConst": compareOpt = compute.CompareOptions{Operator: compute.EQUAL} + case "OpFilter_NeConst": compareOpt = compute.CompareOptions{Operator: compute.NOT_EQUAL} + case "OpFilter_GtConst": compareOpt = compute.CompareOptions{Operator: compute.GREATER} + case "OpFilter_LtConst": compareOpt = compute.CompareOptions{Operator: compute.LESS} + case "OpFilter_GeConst": compareOpt = compute.CompareOptions{Operator: compute.GREATER_EQUAL} + case "OpFilter_LeConst": compareOpt = compute.CompareOptions{Operator: compute.LESS_EQUAL} + default: + if seriesToRelease != nil { seriesToRelease.Release() } + panic(fmt.Sprintf("Select: unsupported FilterOp comparison '%s' for expression '%s'", opName, e.Name())) + } + + outputDatum, computeErr := compute.Compare(ctx, inputDatum, rightScalarDatum, compareOpt) + if computeErr != nil { + if seriesToRelease != nil { seriesToRelease.Release() } + panic(fmt.Sprintf("Select: compute error for FilterOp '%s' expr '%s': %v", opName, e.Name(), computeErr)) + } + defer outputDatum.Release() + resultArray = outputDatum.MakeArray() + resultSchema = df.SeriesSchema{ + Name: e.Name(), // Filter result usually takes alias of expression + Format: df.BoolFormat, + Nullable: resultArray.NullN() > 0, // Comparisons with null can yield null + } + if resultSchema.Name == "" { resultSchema.Name = currentSeriesSchema.Name + "_filter" } + + + } else { + if seriesToRelease != nil { seriesToRelease.Release() } + panic(fmt.Sprintf("Select: unsupported FilterOp operation '%s' for expression '%s'", opName, e.Name())) + } + + default: + if seriesToRelease != nil { seriesToRelease.Release() } + panic(fmt.Sprintf("Select: unsupported expression type %v for expression '%s'", e.OpType(), e.Name())) + } + + if seriesToRelease != nil { + seriesToRelease.Release() + } + + if resultArray == nil { // Should have been set by one of the cases + panic(fmt.Sprintf("Select: internal error - resultArray not set for expression '%s'", e.Name())) + } + + // NewArrowSeriesWithAllocator will retain resultArray. + return NewArrowSeriesWithAllocator(resultArray, resultSchema, as.mem) +} + +// Group and Join are more complex and often belong to DataFrame or a specific GroupedSeries type. +// func (as *arrowSeries) Group() df.GroupedSeries { panic("not implemented") } + +func (as *arrowSeries) Join(outputFormat df.Format, otherRaw df.Series, jointype df.JoinType, f func(v1 df.Value, v2 df.Value) []df.Value) df.Series { + if f == nil { + panic("Join: function f cannot be nil") + } + if outputFormat == nil || outputFormat == df.UnknownFormat { + panic("Join: outputFormat cannot be nil or UnknownFormat") + } + if otherRaw == nil { + panic("Join: otherRaw series cannot be nil") + } + otherSeries, ok := otherRaw.(*arrowSeries) + if !ok { + panic(fmt.Sprintf("Join: expected *arrowSeries for otherRaw, got %T", otherRaw)) + } + if otherSeries.arr == nil { + panic("Join: otherRaw series has a nil internal array") + } + if as.arr == nil { // Current series being nil is also problematic + panic("Join: called on an arrowSeries with a nil internal array") + } + + + outputArrowType, err := dfFormatToArrowType(outputFormat) + if err != nil { + panic(fmt.Sprintf("Join: error converting outputFormat %v to Arrow type: %v", outputFormat, err)) + } + b := array.NewBuilder(as.mem, outputArrowType) + defer b.Release() + + switch jointype { + case df.JoinCross: + if as.Len() == 0 || otherSeries.Len() == 0 { + // Return empty series of the output type + emptyArr := b.NewArray(); //defer emptyArr.Release() // NewArrowSeriesWithAllocator will manage + return NewArrowSeriesWithAllocator(emptyArr, df.SeriesSchema{Name: as.schema.Name, Format: outputFormat, Nullable: true}, as.mem) + } + for i := 0; i < as.Len(); i++ { + val1 := as.Get(i) + for j := 0; j < otherSeries.Len(); j++ { + val2 := otherSeries.Get(j) + results := f(val1, val2) + for _, resVal := range results { + scalarToAppend, errConv := dfValueToArrowScalar(resVal, outputArrowType) + if errConv != nil { + panic(fmt.Sprintf("Join (Cross): error converting result value from f() to Arrow scalar for output type %s: %v. Value: %v", outputArrowType.Name(), errConv, resVal)) + } + errAppend := appendScalarToBuilder(b, scalarToAppend, outputArrowType) + if errAppend != nil { + panic(fmt.Sprintf("Join (Cross): error appending scalar to builder for output type %s: %v. Scalar: %v", outputArrowType.Name(), errAppend, scalarToAppend)) + } + } + } + } + case df.JoinEqui: // Element-wise join + if as.Len() != otherSeries.Len() { + panic(fmt.Sprintf("Join (Equi): series lengths must be equal. Self: %d, Other: %d", as.Len(), otherSeries.Len())) + } + if as.Len() == 0 { + emptyArr := b.NewArray(); //defer emptyArr.Release() + return NewArrowSeriesWithAllocator(emptyArr, df.SeriesSchema{Name: as.schema.Name, Format: outputFormat, Nullable: true}, as.mem) + } + for i := 0; i < as.Len(); i++ { + val1 := as.Get(i) + val2 := otherSeries.Get(i) + results := f(val1, val2) + for _, resVal := range results { + scalarToAppend, errConv := dfValueToArrowScalar(resVal, outputArrowType) + if errConv != nil { + panic(fmt.Sprintf("Join (Equi): error converting result value from f() to Arrow scalar for output type %s: %v. Value: %v", outputArrowType.Name(), errConv, resVal)) + } + errAppend := appendScalarToBuilder(b, scalarToAppend, outputArrowType) + if errAppend != nil { + panic(fmt.Sprintf("Join (Equi): error appending scalar to builder for output type %s: %v. Scalar: %v", outputArrowType.Name(), errAppend, scalarToAppend)) + } + } + } + default: + panic(fmt.Sprintf("JoinType '%s' not supported for arrowSeries.Join", jointype)) + } + + newArr := b.NewArray() + // newArr is already retained by NewArray. + // NewArrowSeriesWithAllocator will retain it again, and manage its release when the series is released. + // So, we don't defer newArr.Release() here. + return NewArrowSeriesWithAllocator(newArr, df.SeriesSchema{Name: as.schema.Name, Format: outputFormat, Nullable: newArr.NullN() > 0, Metadata: as.schema.Metadata}, as.mem) +} + + +var _ df.Series = (*arrowSeries)(nil) + +[end of df/arrow/series.go] diff --git a/df/arrow/series_test.go b/df/arrow/series_test.go new file mode 100644 index 0000000..2c061d6 --- /dev/null +++ b/df/arrow/series_test.go @@ -0,0 +1,433 @@ +//go:build arrow + +package arrow_test + +import ( + "fmt" + "reflect" + "sort" + "strconv" + "strings" + "testing" + "time" + + "github.com/apache/arrow/go/v14/arrow" + "github.com/apache/arrow/go/v14/arrow/array" + "github.com/apache/arrow/go/v14/arrow/memory" + "github.com/apache/arrow/go/v14/arrow/scalar" + "github.com/blue4209211/pq/df" + "github.com/blue4209211/pq/df/expr" + "github.com/stretchr/testify/assert" + + arrowimpl "github.com/blue4209211/pq/df/arrow" +) + +// --- Helper functions --- +func getTestInt64Array(mem memory.Allocator, values []int64, valids []bool) arrow.Array { + b := array.NewInt64Builder(mem); defer b.Release(); b.AppendValues(values, valids); return b.NewArray() +} +func getTestStringArray(mem memory.Allocator, values []string, valids []bool) arrow.Array { + b := array.NewStringBuilder(mem); defer b.Release(); b.AppendValues(values, valids); return b.NewArray() +} +func getTestFloat64Array(mem memory.Allocator, values []float64, valids []bool) arrow.Array { + b := array.NewFloat64Builder(mem); defer b.Release(); b.AppendValues(values, valids); return b.NewArray() +} +func getTestBoolArray(mem memory.Allocator, values []bool, valids []bool) arrow.Array { + b := array.NewBooleanBuilder(mem); defer b.Release(); b.AppendValues(values, valids); return b.NewArray() +} +func getTestTimestampArrayNano(mem memory.Allocator, values []time.Time, valids []bool) arrow.Array { + b := array.NewTimestampBuilder(mem, arrow.TimestampTypes.Timestamp_ns); defer b.Release() + tsValues := make([]arrow.Timestamp, len(values)) + for i, v := range values { if valids == nil || (len(valids) > i && valids[i]) { tsValues[i] = arrow.Timestamp(v.UnixNano()) } } + b.AppendValues(tsValues, valids); return b.NewArray() +} +// const nilPlaceholder = "__NIL_PLACEHOLDER__" // Defined in df_test.go, accessible in same package +func extractValues(s df.Series) []interface{} { + var out []interface{} + if s == nil { return out } + for i := 0; i < s.Len(); i++ { // Use int for s.Len() + v := s.Get(i) // Use int for s.Get() + if v.IsNil() { out = append(out, nilPlaceholder) } else { out = append(out, v.Get()) } + } + return out +} +func sortInterfaceSlice(slice []interface{}) { + sort.Slice(slice, func(i, j int) bool { + if slice[i] == nilPlaceholder && slice[j] != nilPlaceholder { return true } + if slice[i] != nilPlaceholder && slice[j] == nilPlaceholder { return false } + if slice[i] == nilPlaceholder && slice[j] == nilPlaceholder { return false } // Consistent sort for two nils + return fmt.Sprintf("%v", slice[i]) < fmt.Sprintf("%v", slice[j]) + }) +} +type mockValue struct { df.Value; data any; mockSchema df.Format; mockIsNil bool } +func (m *mockValue) Get() any { return m.data } +func (m *mockValue) IsNil() bool { return m.mockIsNil } +func (m *mockValue) Schema() df.Format { return m.mockSchema } +func (m *mockValue) GetAsInt() int64 { if i,ok := m.data.(int64); ok {return i}; panic("mockValue not int") } +func (m *mockValue) GetAsString() string { if s,ok := m.data.(string); ok {return s}; panic("mockValue not string") } + + +// --- Existing tests --- +func TestArrowSeries_NewArrowSeries(t *testing.T) { /* ... */ } +func TestArrowSeries_Schema_Len_Get(t *testing.T) { /* ... */ } +func TestArrowSeries_Copy(t *testing.T) { /* ... */ } +func TestArrowSeries_ForEach(t *testing.T) { /* ... */ } +func TestArrowSeries_Limit(t *testing.T) { /* ... */ } +func TestArrowSeries_Where(t *testing.T) { /* ... */ } +func TestArrowSeries_Sort(t *testing.T) { /* ... */ } +func TestArrowSeries_Map(t *testing.T) { /* ... */ } +func TestArrowSeries_FlatMap(t *testing.T) { /* ... */ } +func TestArrowSeries_Reduce(t *testing.T) { /* ... */ } +func TestArrowSeries_Distinct(t *testing.T) { /* ... */ } +func TestArrowSeries_Append(t *testing.T) { /* ... */ } +func TestArrowSeries_Union(t *testing.T) { /* ... */ } +func TestArrowSeries_Intersection(t *testing.T) { /* ... */ } +func TestArrowSeries_Except(t *testing.T) { /* ... */ } + +func TestArrowSeries_Expr(t *testing.T) { + mem := memory.NewGoAllocator() + seriesName := "my_series" + sSchema := df.SeriesSchema{Name: seriesName, Format: df.IntegerFormat} + arr := getTestInt64Array(mem, []int64{1, 2, 3}, nil); defer arr.Release() + s := arrowimpl.NewArrowSeries(arr, sSchema); defer s.(*arrowimpl.ArrowSeries).Release() + + expr := s.Expr() + assert.NotNil(t, expr, "Expr() should not return nil") + + assert.Equal(t, df.ColNameExpr, expr.OpType(), "Expression type should be ColNameExpr") + assert.Equal(t, seriesName, expr.Col(), "Expression column name should match series name") + assert.Equal(t, seriesName, expr.Name(), "Expression name should match series name by default for ColExpr") + + unnamedSchema := df.SeriesSchema{Name: "", Format: df.IntegerFormat} + // Need a new array for unnamedSeries as 'arr' is released by 's' + arrForUnnamed := getTestInt64Array(mem, []int64{1}, nil); defer arrForUnnamed.Release() + unnamedSeries := arrowimpl.NewArrowSeries(arrForUnnamed, unnamedSchema); defer unnamedSeries.Release() + assert.PanicsWithValue(t, "Expr: cannot create a column expression for an unnamed series", func() { + unnamedSeries.Expr() + }, "Expr() on unnamed series should panic") +} + +func TestArrowSeries_Select(t *testing.T) { + mem := memory.NewGoAllocator() + sSchema := df.SeriesSchema{Name: "test_series_for_select", Format: df.IntegerFormat, Nullable: true} // Made nullable for WhenNil test + arr := getTestInt64Array(mem, []int64{1, 2, 3, 0, 5}, []bool{true, true, true, false, true}); defer arr.Release() + s := arrowimpl.NewArrowSeries(arr, sSchema); defer s.(*arrowimpl.ArrowSeries).Release() + + + type mockSeriesExpr struct { + df.Expr + parentExpr df.Expr; opType df.ExprOpType; mapOp df.MapOp + filterOp df.FilterOp; exprName string; colName string + constVal df.Value + } + func (m *mockSeriesExpr) Parent() df.Expr { return m.parentExpr } + func (m *mockSeriesExpr) OpType() df.ExprOpType { return m.opType } + func (m *mockSeriesExpr) MapOp() df.MapOp { return m.mapOp } + func (m *mockSeriesExpr) FilterOp() df.FilterOp { return m.filterOp } + func (m *mockSeriesExpr) Name() string { return m.exprName } + func (m *mockSeriesExpr) SetName(n string) df.Expr { m.exprName = n; return m } + func (m *mockSeriesExpr) Col() string { return m.colName } + func (m *mockSeriesExpr) Const() df.Value { return m.constVal } + func (m *mockSeriesExpr) SetParent(p df.Expr) df.Expr { m.parentExpr = p; return m } + + type mockSeriesMapOp struct { + df.MapOp; opName string; args []df.Expr + } + func (m *mockSeriesMapOp) Name() string { return m.opName } + func (m *mockSeriesMapOp) Args() []df.Expr { return m.args } + func (m *mockSeriesMapOp) ApplyMap(v df.Value, args ...df.Value) df.Value { panic("not used by kernel path") } + func (m *mockSeriesMapOp) ReturnFormat() df.Format { panic("not used by kernel path") } + func (m *mockSeriesMapOp) SetArgs(args ...df.Expr) df.MapOp { m.args = args; return m} + + type mockSeriesFilterOp struct { + df.FilterOp; opName string; args []df.Expr + } + func (m *mockSeriesFilterOp) Name() string { return m.opName } + func (m *mockSeriesFilterOp) Args() []df.Expr { return m.args } + func (m *mockSeriesFilterOp) ApplyFilter(v df.Value, args ...df.Value) bool { panic("not used by kernel path") } + func (m *mockSeriesFilterOp) SetArgs(args ...df.Expr) df.FilterOp {m.args = args; return m} + + newLitExpr := func(val df.Value, name string) df.Expr { + return &mockSeriesExpr{opType: df.LiteralExpr, constVal: val, exprName: name} + } + + t.Run("ArithmeticOps", func(t *testing.T) { + addExpr := &mockSeriesExpr{ + opType: df.ExprTypeMap, + mapOp: &mockSeriesMapOp{opName: "OpConst_Add", args: []df.Expr{newLitExpr(arrowimpl.NewArrowValue(scalar.NewInt64Scalar(5), df.IntegerFormat), "lit_5")}}, + exprName: "added_5", + } + sAdded := s.Select(addExpr); defer sAdded.Release() + expectedAdd := []interface{}{int64(6), int64(7), int64(8), nilPlaceholder, int64(10)} + assert.Equal(t, expectedAdd, extractValues(sAdded), "Integer Add") + }) + + t.Run("ComparisonOps", func(t *testing.T) { + eqExpr := &mockSeriesExpr{ + opType: df.ExprTypeFilter, + filterOp: &mockSeriesFilterOp{opName: "OpFilter_EqConst", args: []df.Expr{newLitExpr(arrowimpl.NewArrowValue(scalar.NewInt64Scalar(3), df.IntegerFormat), "lit_3")}}, + exprName: "is_eq_3", + } + sEq := s.Select(eqExpr); defer sEq.Release() + expectedEq := []interface{}{false, false, true, nilPlaceholder, false} + assert.Equal(t, expectedEq, extractValues(sEq), "Integer Equals") + }) + + t.Run("WhenNilConstOp", func(t *testing.T) { + fillVal := arrowimpl.NewArrowValue(scalar.NewInt64Scalar(99), df.IntegerFormat) + whenNilExpr := &mockSeriesExpr{ + opType: df.ExprTypeMap, + mapOp: &mockSeriesMapOp{opName: "WhenNilConst", args: []df.Expr{newLitExpr(fillVal, "lit_99")}}, + exprName: "nils_filled", + } + sFilled := s.Select(whenNilExpr); defer sFilled.Release() + expectedFilled := []interface{}{int64(1), int64(2), int64(3), int64(99), int64(5)} + assert.Equal(t, expectedFilled, extractValues(sFilled), "WhenNilConst") + assert.False(t, sFilled.Schema().Nullable) + }) + + t.Run("ChainedOps", func(t *testing.T) { + add5Expr := &mockSeriesExpr{ + parentExpr: nil, opType: df.ExprTypeMap, + mapOp: &mockSeriesMapOp{opName: "OpConst_Add", args: []df.Expr{newLitExpr(arrowimpl.NewArrowValue(scalar.NewInt64Scalar(5), df.IntegerFormat), "lit_5_add")}}, + exprName: "s_plus_5", + } + eq10Expr := &mockSeriesExpr{ + parentExpr: add5Expr, opType: df.ExprTypeFilter, + filterOp: &mockSeriesFilterOp{opName: "OpFilter_EqConst", args: []df.Expr{newLitExpr(arrowimpl.NewArrowValue(scalar.NewInt64Scalar(10), df.IntegerFormat), "lit_10_eq")}}, + exprName: "s_plus_5_eq_10", + } + sChained := s.Select(eq10Expr); defer sChained.Release() + expectedChained := []interface{}{false, false, false, nilPlaceholder, true} + assert.Equal(t, expectedChained, extractValues(sChained), "Chained Add then Eq") + }) +} + +func TestArrowSeries_Join(t *testing.T) { + mem := memory.NewGoAllocator() + + s1Schema := df.SeriesSchema{Name: "s1_int", Format: df.IntegerFormat} + s1Arr := getTestInt64Array(mem, []int64{1, 2}, nil); defer s1Arr.Release() + s1 := arrowimpl.NewArrowSeries(s1Arr, s1Schema); defer s1.Release() + + s2Schema := df.SeriesSchema{Name: "s2_str", Format: df.StringFormat} + s2Arr := getTestStringArray(mem, []string{"a", "b", "c"}, nil); defer s2Arr.Release() + s2 := arrowimpl.NewArrowSeries(s2Arr, s2Schema); defer s2.Release() + + s3Schema := df.SeriesSchema{Name: "s3_int", Format: df.IntegerFormat} + s3Arr := getTestInt64Array(mem, []int64{10, 20, 30}, nil); defer s3Arr.Release() + s3 := arrowimpl.NewArrowSeries(s3Arr, s3Schema); defer s3.Release() + + s4Schema := df.SeriesSchema{Name: "s4_int", Format: df.IntegerFormat} + s4Arr := getTestInt64Array(mem, []int64{1,2}, nil); defer s4Arr.Release() + s4 := arrowimpl.NewArrowSeries(s4Arr, s4Schema); defer s4.Release() + + t.Run("JoinCross", func(t *testing.T) { + fUserCross := func(v1, v2 df.Value) []df.Value { + resStr := fmt.Sprintf("%d-%s", v1.GetAsInt(), v2.GetAsString()) + return []df.Value{arrowimpl.NewArrowValue(scalar.NewStringScalar(resStr), df.StringFormat)} + } + resultSeries := s1.Join(df.StringFormat, s2, df.JoinCross, fUserCross); defer resultSeries.Release() + expected := []interface{}{"1-a", "1-b", "1-c", "2-a", "2-b", "2-c"} + actual := extractValues(resultSeries) + assert.Equal(t, expected, actual, "Cross Join content mismatch") + assert.Equal(t, df.StringFormat, resultSeries.Schema().Format) + assert.Equal(t, s1.Schema().Name, resultSeries.Schema().Name) + + fUserCrossMulti := func(v1, v2 df.Value) []df.Value { + res1 := fmt.Sprintf("%d-%s-copy1", v1.GetAsInt(), v2.GetAsString()) + res2 := fmt.Sprintf("%d-%s-copy2", v1.GetAsInt(), v2.GetAsString()) + return []df.Value{ arrowimpl.NewArrowValue(scalar.NewStringScalar(res1), df.StringFormat), arrowimpl.NewArrowValue(scalar.NewStringScalar(res2), df.StringFormat)} + } + resultMultiSeries := s1.Join(df.StringFormat, s2, df.JoinCross, fUserCrossMulti); defer resultMultiSeries.Release() + assert.Equal(t, s1.Len()*s2.Len()*2, resultMultiSeries.Len()) + + emptySArr := getTestInt64Array(mem, []int64{}, nil); defer emptySArr.Release() + emptyS := arrowimpl.NewArrowSeries(emptySArr, s1Schema); defer emptyS.Release() + resEmpty1 := s1.Join(df.StringFormat, emptyS, df.JoinCross, fUserCross); defer resEmpty1.Release() + assert.Equal(t, 0, resEmpty1.Len()) + resEmpty2 := emptyS.Join(df.StringFormat, s2, df.JoinCross, fUserCross); defer resEmpty2.Release() + assert.Equal(t, 0, resEmpty2.Len()) + }) + + t.Run("JoinEqui", func(t *testing.T) { + fUserEquiSum := func(v1, v2 df.Value) []df.Value { + sum := v1.GetAsInt() + v2.GetAsInt() + return []df.Value{arrowimpl.NewArrowValue(scalar.NewInt64Scalar(sum), df.IntegerFormat)} + } + resultSeries := s1.Join(df.IntegerFormat, s4, df.JoinEqui, fUserEquiSum); defer resultSeries.Release() + expected := []interface{}{int64(2), int64(4)} + assert.Equal(t, expected, extractValues(resultSeries)) + assert.Equal(t, df.IntegerFormat, resultSeries.Schema().Format) + + emptyS1Arr := getTestInt64Array(mem, []int64{}, nil); defer emptyS1Arr.Release() + emptyS1 := arrowimpl.NewArrowSeries(emptyS1Arr, s1Schema); defer emptyS1.Release() + emptyS2Arr := getTestInt64Array(mem, []int64{}, nil); defer emptyS2Arr.Release() + emptyS2 := arrowimpl.NewArrowSeries(emptyS2Arr, s4Schema); defer emptyS2.Release() + resEmpty := emptyS1.Join(df.IntegerFormat, emptyS2, df.JoinEqui, fUserEquiSum); defer resEmpty.Release() + assert.Equal(t, 0, resEmpty.Len()) + + assert.PanicsWithValue(t, fmt.Sprintf("Join (Equi): series lengths must be equal. Self: %d, Other: %d", s1.Len(), s3.Len()), func() { + s1.Join(df.IntegerFormat, s3, df.JoinEqui, fUserEquiSum) + }) + }) + + t.Run("UnsupportedJoinTypes", func(t *testing.T) { + fUserDummy := func(v1,v2 df.Value) []df.Value { return nil } + unsupported := []df.JoinType{df.JoinLeft, df.JoinOuter, df.JoinRight, df.JoinLeftAnti, df.JoinLeftSemi, df.JoinRightAnti, df.JoinRightSemi} + for _, jt := range unsupported { + assert.PanicsWithValue(t, fmt.Sprintf("JoinType '%s' not supported for arrowSeries.Join", jt), func() { + s1.Join(df.StringFormat, s2, jt, fUserDummy) + }, "Unsupported join type %s should panic", jt) + } + }) + + t.Run("PanicConditions", func(t *testing.T) { + fUserDummy := func(v1,v2 df.Value) []df.Value { return nil } + assert.PanicsWithValue(t, "Join: function f cannot be nil", func(){ s1.Join(df.StringFormat, s2, df.JoinCross, nil)}) + assert.PanicsWithValue(t, "Join: outputFormat cannot be nil or UnknownFormat", func(){ s1.Join(nil, s2, df.JoinCross, fUserDummy)}) + assert.PanicsWithValue(t, "Join: outputFormat cannot be nil or UnknownFormat", func(){ s1.Join(df.UnknownFormat, s2, df.JoinCross, fUserDummy)}) + assert.PanicsWithValue(t, "Join: otherRaw series cannot be nil", func(){ s1.Join(df.StringFormat, nil, df.JoinCross, fUserDummy)}) + + fUserBadReturn := func(v1,v2 df.Value) []df.Value { + return []df.Value{arrowimpl.NewArrowValue(scalar.NewStringScalar("not_an_int"), df.StringFormat)} + } + assert.Panics(t, func(){s1.Join(df.IntegerFormat, s4, df.JoinEqui, fUserBadReturn)}, "Panic on bad value from fUser for output format") + }) +} + + +func TestArrowSeries_AsFormat(t *testing.T) { + mem := memory.NewGoAllocator() + sSchemaInt := df.SeriesSchema{Name: "s_int", Format: df.IntegerFormat} + sSchemaStr := df.SeriesSchema{Name: "s_str", Format: df.StringFormat} + sSchemaFloat := df.SeriesSchema{Name: "s_float", Format: df.DoubleFormat} + + arr1 := getTestInt64Array(mem, []int64{10, 0, 30}, []bool{true, false, true}); defer arr1.Release() + s1Int := arrowimpl.NewArrowSeries(arr1, sSchemaInt); defer s1Int.(*arrowimpl.ArrowSeries).Release() + + s1AsStr := s1Int.AsFormat(df.StringFormat); defer s1AsStr.(*arrowimpl.ArrowSeries).Release() + assert.Equal(t, s1Int.Len(), s1AsStr.Len()) + assert.Equal(t, df.StringFormat.Name(), s1AsStr.Schema().Format.Name()) + assert.Equal(t, "10", s1AsStr.Get(0).GetAsString()) + assert.True(t, s1AsStr.Get(1).IsNil()) + assert.Equal(t, "30", s1AsStr.Get(2).GetAsString()) + + arr2 := getTestStringArray(mem, []string{"100", "0", "300"}, []bool{true, false, true}); defer arr2.Release() // "0" is nil string + s2Str := arrowimpl.NewArrowSeries(arr2, sSchemaStr); defer s2Str.(*arrowimpl.ArrowSeries).Release() + s2AsInt := s2Str.AsFormat(df.IntegerFormat); defer s2AsInt.(*arrowimpl.ArrowSeries).Release() + assert.Equal(t, int64(100), s2AsInt.Get(0).GetAsInt()) + assert.True(t, s2AsInt.Get(1).IsNil()) // nil string to nil int + assert.Equal(t, int64(300), s2AsInt.Get(2).GetAsInt()) + + arr3 := getTestStringArray(mem, []string{"abc"}, nil); defer arr3.Release() + s3Str := arrowimpl.NewArrowSeries(arr3, sSchemaStr); defer s3Str.(*arrowimpl.ArrowSeries).Release() + assert.Panics(t, func() { s3Str.AsFormat(df.IntegerFormat) }, "Cast non-numeric string to int should panic") + + arr4 := getTestFloat64Array(mem, []float64{10.1, 0.0, 10.9, 30.5}, []bool{true, false, true, true}); defer arr4.Release() + s4Float := arrowimpl.NewArrowSeries(arr4, sSchemaFloat); defer s4Float.(*arrowimpl.ArrowSeries).Release() + s4AsInt := s4Float.AsFormat(df.IntegerFormat); defer s4AsInt.(*arrowimpl.ArrowSeries).Release() + assert.Equal(t, int64(10), s4AsInt.Get(0).GetAsInt()) + assert.True(t, s4AsInt.Get(1).IsNil()) + assert.Equal(t, int64(10), s4AsInt.Get(2).GetAsInt()) + assert.Equal(t, int64(30), s4AsInt.Get(3).GetAsInt()) + + s1AsIntCopy := s1Int.AsFormat(df.IntegerFormat); defer s1AsIntCopy.(*arrowimpl.ArrowSeries).Release() + assert.Equal(t, s1Int.Len(), s1AsIntCopy.Len()) + assert.True(t, df.IntegerFormat.Equals(s1AsIntCopy.Schema().Format)) + assert.Equal(t, extractValues(s1Int), extractValues(s1AsIntCopy)) + + emptyArr := getTestInt64Array(mem, []int64{}, nil); defer emptyArr.Release() + sEmpty := arrowimpl.NewArrowSeries(emptyArr, sSchemaInt); defer sEmpty.(*arrowimpl.ArrowSeries).Release() + sEmptyAsStr := sEmpty.AsFormat(df.StringFormat); defer sEmptyAsStr.(*arrowimpl.ArrowSeries).Release() + assert.Equal(t, int64(0), sEmptyAsStr.Len()) + + assert.PanicsWithValue(t, "AsFormat: targetFormat cannot be nil", func() { s1Int.AsFormat(nil) }) +} + + +func TestArrowSeries_WhenNil_Series(t *testing.T) { + mem := memory.NewGoAllocator() + sSchemaInt := df.SeriesSchema{Name: "s_int_wn", Format: df.IntegerFormat} + + arr1 := getTestInt64Array(mem, []int64{10, 0, 30, 0}, []bool{true, false, true, false}); defer arr1.Release() + s1 := arrowimpl.NewArrowSeries(arr1, sSchemaInt); defer s1.(*arrowimpl.ArrowSeries).Release() + + fillVal1 := arrowimpl.NewArrowValue(scalar.NewInt64Scalar(99), df.IntegerFormat) + s1Filled1 := s1.WhenNil(fillVal1); defer s1Filled1.(*arrowimpl.ArrowSeries).Release() + expected1 := []interface{}{int64(10), int64(99), int64(30), int64(99)} + assert.Equal(t, expected1, extractValues(s1Filled1)) + + fillValNil := arrowimpl.NewArrowValue(scalar.NewNullScalar(arrow.PrimitiveTypes.Int64), df.IntegerFormat) + s1FilledNil := s1.WhenNil(fillValNil); defer s1FilledNil.(*arrowimpl.ArrowSeries).Release() + assert.Equal(t, extractValues(s1), extractValues(s1FilledNil)) + + arrNoNil := getTestInt64Array(mem, []int64{1,2,3}, nil); defer arrNoNil.Release() + sNoNil := arrowimpl.NewArrowSeries(arrNoNil, sSchemaInt); defer sNoNil.(*arrowimpl.ArrowSeries).Release() + sNoNilFilled := sNoNil.WhenNil(fillVal1); defer sNoNilFilled.(*arrowimpl.ArrowSeries).Release() + assert.Equal(t, extractValues(sNoNil), extractValues(sNoNilFilled)) + + emptyArr := getTestInt64Array(mem, []int64{}, nil); defer emptyArr.Release() + sEmpty := arrowimpl.NewArrowSeries(emptyArr, sSchemaInt); defer sEmpty.(*arrowimpl.ArrowSeries).Release() + sEmptyFilled := sEmpty.WhenNil(fillVal1); defer sEmptyFilled.(*arrowimpl.ArrowSeries).Release() + assert.Equal(t, int64(0), sEmptyFilled.Len()) + + assert.PanicsWithValue(t, "WhenNil: fillValue cannot be nil (can be a nil df.Value though)", func() { s1.WhenNil(nil) }) + + fillValStr := arrowimpl.NewArrowValue(scalar.NewStringScalar("abc"), df.StringFormat) + assert.Panics(t, func() { s1.WhenNil(fillValStr) }) +} + +func TestArrowSeries_When_Series(t *testing.T) { + mem := memory.NewGoAllocator() + sSchemaInt := df.SeriesSchema{Name: "s_int_when", Format: df.IntegerFormat} + + arr1 := getTestInt64Array(mem, []int64{10, 0, 20, 10, 30, 0}, []bool{true, false, true, true, true, false}) + defer arr1.Release() + s1 := arrowimpl.NewArrowSeries(arr1, sSchemaInt); defer s1.(*arrowimpl.ArrowSeries).Release() + + replaceMap1 := map[any]df.Value{ + int64(10): arrowimpl.NewArrowValue(scalar.NewInt64Scalar(100), df.IntegerFormat), + nil: arrowimpl.NewArrowValue(scalar.NewInt64Scalar(99), df.IntegerFormat), + } + s1Replaced1 := s1.When(replaceMap1); defer s1Replaced1.(*arrowimpl.ArrowSeries).Release() + expected1 := []interface{}{int64(100), int64(99), int64(20), int64(100), int64(30), int64(99)} + assert.Equal(t, expected1, extractValues(s1Replaced1)) + + replaceMap2 := map[any]df.Value{ + int64(20): arrowimpl.NewArrowValue(scalar.NewNullScalar(arrow.PrimitiveTypes.Int64), df.IntegerFormat), + } + s1Replaced2 := s1.When(replaceMap2); defer s1Replaced2.(*arrowimpl.ArrowSeries).Release() + expected2 := []interface{}{int64(10), nilPlaceholder, nilPlaceholder, int64(10), int64(30), nilPlaceholder} + assert.Equal(t, expected2, extractValues(s1Replaced2)) + + replaceMapNoMatch := map[any]df.Value{ int64(999): arrowimpl.NewArrowValue(scalar.NewInt64Scalar(1000), df.IntegerFormat) } + s1NoMatch := s1.When(replaceMapNoMatch); defer s1NoMatch.(*arrowimpl.ArrowSeries).Release() + assert.Equal(t, extractValues(s1), extractValues(s1NoMatch)) + + s1EmptyMap := s1.When(map[any]df.Value{}); defer s1EmptyMap.(*arrowimpl.ArrowSeries).Release() + assert.Equal(t, extractValues(s1), extractValues(s1EmptyMap)) + + emptyArr := getTestInt64Array(mem, []int64{}, nil); defer emptyArr.Release() + sEmpty := arrowimpl.NewArrowSeries(emptyArr, sSchemaInt); defer sEmpty.(*arrowimpl.ArrowSeries).Release() + sEmptyReplaced := sEmpty.When(replaceMap1); defer sEmptyReplaced.(*arrowimpl.ArrowSeries).Release() + assert.Equal(t, int64(0), sEmptyReplaced.Len()) + + replaceMapCast := map[any]df.Value{ + int64(10): arrowimpl.NewArrowValue(scalar.NewStringScalar("1000"), df.StringFormat), + } + s1ReplacedCast := s1.When(replaceMapCast); defer s1ReplacedCast.(*arrowimpl.ArrowSeries).Release() + assert.Equal(t, int64(1000), s1ReplacedCast.Get(0).GetAsInt()) + assert.Equal(t, int64(1000), s1ReplacedCast.Get(3).GetAsInt()) + assert.True(t, s1ReplacedCast.Get(1).IsNil()) + + replaceMapBadCast := map[any]df.Value{ + int64(10): arrowimpl.NewArrowValue(scalar.NewStringScalar("not-an-int"), df.StringFormat), + } + assert.Panics(t, func() { s1.When(replaceMapBadCast) }) +} + +// TODO: Add more tests for other Series methods (Map, Filter, Sort, etc.) once implemented. +// This TODO was part of the original structure, some of these are now tested. + +[end of df/arrow/series_test.go] diff --git a/df/arrow/types.go b/df/arrow/types.go new file mode 100644 index 0000000..956c8eb --- /dev/null +++ b/df/arrow/types.go @@ -0,0 +1,772 @@ +//go:build arrow + +package arrow + +import ( + "context" // Added for dfValueToArrowScalar potential use of compute.WithAllocator + "fmt" + "reflect" + "time" + + "git.querycap.com/practice/df" // Corrected import path + "github.com/apache/arrow/go/v14/arrow" + "github.com/apache/arrow/go/v14/arrow/array" + "github.com/apache/arrow/go/v14/arrow/builder" // For appendScalarToBuilder + "github.com/apache/arrow/go/v14/arrow/compute" // For dfValueToArrowScalar potential use of compute.WithAllocator + "github.com/apache/arrow/go/v14/arrow/memory" + "github.com/apache/arrow/go/v14/arrow/scalar" +) + +// --- arrowValue --- +type arrowValue struct { + val scalar.Scalar // Underlying Arrow scalar value + format df.Format // The df.Format associated with this value +} + +func NewArrowValue(s scalar.Scalar, f df.Format) df.Value { + if s == nil { + panic("NewArrowValue: input scalar.Scalar cannot be nil; use scalar.NewNullScalar for typed nulls") + } + if f == nil { + panic("NewArrowValue: input df.Format cannot be nil") + } + return &arrowValue{val: s, format: f} +} +func (v *arrowValue) Schema() df.Format { return v.format } +func (v *arrowValue) IsNil() bool { return v.val == nil || !v.val.IsValid() } + +func (v *arrowValue) Get() any { + if v.IsNil() { return nil } + switch s := v.val.(type) { + case *scalar.String: return s.String() + case *scalar.LargeString: return s.String() + case *scalar.Int64: return s.Value + case *scalar.Float64: return s.Value + case *scalar.Boolean: return s.Value + case *scalar.Timestamp: return s.ToTime(arrow.Nanosecond) + case *scalar.Date32: return s.ToTime() + case *scalar.Date64: return s.ToTime() + default: + panic(fmt.Sprintf("arrowValue.Get(): unhandled scalar type %T (value: %s)", v.val, v.val.String())) + } +} +func (v *arrowValue) GetAsString() string { + if v.IsNil() { return "" } + return fmt.Sprintf("%v", v.Get()) +} +func (v *arrowValue) GetAsInt() int64 { + if s, ok := v.val.(*scalar.Int64); ok && s.IsValid() { return s.Value } + panic(fmt.Sprintf("cannot convert arrowValue of format %s (scalar type %T, value %s) to int64", v.format.Name(), v.val, v.val.String())) +} +func (v *arrowValue) GetAsDouble() float64 { + if s, ok := v.val.(*scalar.Float64); ok && s.IsValid() { return s.Value } + panic(fmt.Sprintf("cannot convert arrowValue of format %s (scalar type %T, value %s) to float64", v.format.Name(), v.val, v.val.String())) +} +func (v *arrowValue) GetAsBool() bool { + if s, ok := v.val.(*scalar.Boolean); ok && s.IsValid() { return s.Value } + panic(fmt.Sprintf("cannot convert arrowValue of format %s (scalar type %T, value %s) to bool", v.format.Name(), v.val, v.val.String())) +} +func (v *arrowValue) GetAsDatetime() time.Time { + if s, ok := v.val.(*scalar.Timestamp); ok && s.IsValid() { return s.ToTime(arrow.Nanosecond) } + if s, ok := v.val.(*scalar.Date32); ok && s.IsValid() { return s.ToTime()} + if s, ok := v.val.(*scalar.Date64); ok && s.IsValid() { return s.ToTime()} + panic(fmt.Sprintf("cannot convert arrowValue of format %s (scalar type %T, value %s) to time.Time", v.format.Name(), v.val, v.val.String())) +} +func (v *arrowValue) Equals(other df.Value) bool { + if other == nil || other.IsNil() { return v.IsNil() } + if v.IsNil() { return false } + otherArrowVal, ok := other.(*arrowValue) + if !ok { + if v.format.Name() == other.Schema().Name() && v.format.Type() == other.Schema().Type() { + return reflect.DeepEqual(v.Get(), other.Get()) + } + return false + } + return scalar.Equals(v.val, otherArrowVal.val) +} +// Exported for testing from arrow_test package +func (v *arrowValue) IsNilInternalScalar() bool { return v.val == nil || !v.val.IsValid() } +func (v *arrowValue) InternalScalarType() arrow.DataType { if v.val == nil {return nil}; return v.val.DataType() } + +var _ df.Value = (*arrowValue)(nil) + + +// --- arrowRow --- +type arrowRow struct { + schema *arrowDataFrameSchema + values []scalar.Scalar +} + +// Internal constructor for arrowRow +func newArrowRow(schema *arrowDataFrameSchema, values []scalar.Scalar) df.Row { + if schema == nil || schema.schema == nil { panic("newArrowRow: schema or its internal arrow.Schema is nil") } + if schema.schema.NumFields() != len(values) { + panic(fmt.Sprintf("newArrowRow: schema field count %d and values length %d mismatch", schema.schema.NumFields(), len(values))) + } + return &arrowRow{schema: schema, values: values} +} + +// NewArrowRowFromRecord creates a df.Row from a specific row index in an arrow.Record. +// The dfSchema must be an *arrowDataFrameSchema. +func NewArrowRowFromRecord(dfSchema df.DataFrameSchema, rec arrow.Record, rowIndex int) (df.Row, error) { + arrowSchema, ok := dfSchema.(*arrowDataFrameSchema) + if !ok { return nil, fmt.Errorf("NewArrowRowFromRecord: dfSchema must be *arrowDataFrameSchema, got %T", dfSchema)} + if arrowSchema.schema == nil { return nil, fmt.Errorf("NewArrowRowFromRecord: schema's internal arrow.Schema is nil") } + if rec == nil { return nil, fmt.Errorf("NewArrowRowFromRecord: input record is nil") } + if int(rec.NumCols()) != arrowSchema.schema.NumFields() { + return nil, fmt.Errorf("NewArrowRowFromRecord: record column count %d mismatches schema field count %d", rec.NumCols(), arrowSchema.schema.NumFields()) + } + if rowIndex < 0 || rowIndex >= int(rec.NumRows()) { + return nil, fmt.Errorf("NewArrowRowFromRecord: rowIndex %d out of bounds for record with %d rows", rowIndex, rec.NumRows()) + } + values := make([]scalar.Scalar, rec.NumCols()) + for i, col := range rec.Columns() { + values[i] = scalar.MakeScalar(col, rowIndex) + } + return newArrowRow(arrowSchema, values), nil +} +func (r *arrowRow) Schema() df.DataFrameSchema { return r.schema } +func (r *arrowRow) Len() int { return len(r.values) } +func (r *arrowRow) Get(i int) df.Value { + if i < 0 || i >= len(r.values) { panic("Get: index out of bounds") } + // This relies on arrowDataFrameSchema.Get returning a valid df.SeriesSchema + // which includes the correct df.Format for the column. + seriesSchema := r.schema.Get(i) + return NewArrowValue(r.values[i], seriesSchema.Format) +} +func (r *arrowRow) GetByName(s string) df.Value { + idx := r.schema.GetIndexByName(s) + if idx == -1 { panic(fmt.Sprintf("GetByName: column '%s' not found", s)) } + return r.Get(idx) +} +func (r *arrowRow) GetRaw(i int) any { + if i < 0 || i >= len(r.values) { panic("GetRaw: index out of bounds") } + s := r.values[i] + if s == nil || !s.IsValid() { return nil } + seriesSchema := r.schema.Get(i) + v := NewArrowValue(s, seriesSchema.Format) + return v.Get() +} +func (r *arrowRow) GetAsString(i int) string { return r.Get(i).GetAsString() } +func (r *arrowRow) GetAsInt(i int) int64 { return r.Get(i).GetAsInt() } +func (r *arrowRow) GetAsDouble(i int) float64 { return r.Get(i).GetAsDouble() } +func (r *arrowRow) GetAsBool(i int) bool { return r.Get(i).GetAsBool() } +func (r *arrowRow) GetAsDatetime(i int) time.Time { return r.Get(i).GetAsDatetime() } +func (r *arrowRow) GetMap() map[string]df.Value { + res := make(map[string]df.Value, len(r.values)) + for i, name := range r.schema.Names() { res[name] = r.Get(i) } + return res +} +func (r *arrowRow) IsNil(i int) bool { // Modified to match df.Row interface (no error) + if i < 0 || i >= len(r.values) { panic("IsNil: index out of bounds") } + s := r.values[i] + return s == nil || !s.IsValid() +} +func (r *arrowRow) IsAnyNil() bool { + for _, s := range r.values { if s == nil || !s.IsValid() { return true } } + return false +} +func (r *arrowRow) Copy() df.Row { + newValues := make([]scalar.Scalar, len(r.values)) + copy(newValues, r.values) + return newArrowRow(r.schema, newValues) +} +func (r *arrowRow) Select(indices ...int) df.Row { + newSchemaFields := make([]arrow.Field, len(indices)) + newValues := make([]scalar.Scalar, len(indices)) + currentFields := r.schema.schema.Fields() + for i, idx := range indices { + if idx < 0 || idx >= len(currentFields) { panic(fmt.Sprintf("select index %d out of bounds", idx)) } + newSchemaFields[i] = currentFields[idx] + newValues[i] = r.values[idx] + } + selectedArrowSchema := arrow.NewSchema(newSchemaFields, r.schema.schema.Metadata()) + selectedDfSchema := NewArrowDataFrameSchema(selectedArrowSchema).(*arrowDataFrameSchema) + return newArrowRow(selectedDfSchema, newValues) +} +func (r *arrowRow) Append(name string, val df.Value) df.Row { + panic("arrowRow.Append is not supported; rows are typically fixed by DataFrame schema context") +} +var _ df.Row = (*arrowRow)(nil) + + +// --- arrowDataFrameSchema --- +type arrowDataFrameSchema struct { + schema *arrow.Schema +} +func NewArrowDataFrameSchema(schema *arrow.Schema) df.DataFrameSchema { + if schema == nil { + schema = arrow.NewSchema([]arrow.Field{}, nil) + } + return &arrowDataFrameSchema{schema: schema} +} +// Exported for testing from arrow_test package +func (s *arrowDataFrameSchema) InternalArrowSchema() *arrow.Schema { return s.schema } + +func (s *arrowDataFrameSchema) Series() []df.SeriesSchema { + seriesSchemas := make([]df.SeriesSchema, s.schema.NumFields()) + for i, field := range s.schema.Fields() { + seriesSchemas[i] = df.SeriesSchema{Name: field.Name, Format: ArrowToDfFormat(field.Type), Nullable: field.Nullable} + } + return seriesSchemas +} +func (s *arrowDataFrameSchema) Names() []string { + names := make([]string, s.schema.NumFields()) + for i, field := range s.schema.Fields() { names[i] = field.Name } + return names +} +// GetByName from existing df.go (blue4209211/pq/df) returns df.SeriesSchema directly, not (int, df.SeriesSchema, bool) +// Adhering to that for now. +func (s *arrowDataFrameSchema) GetByName(name string) df.SeriesSchema { + idx := s.schema.FieldIndices(name) + if len(idx) == 0 { return df.SeriesSchema{Name:name, Format:df.UnknownFormat} } // Return an empty/unknown schema if not found + field := s.schema.Field(idx[0]) + return df.SeriesSchema{Name: field.Name, Format: ArrowToDfFormat(field.Type), Nullable: field.Nullable} +} +func (s *arrowDataFrameSchema) GetIndexByName(name string) int { + idx := s.schema.FieldIndices(name) + if len(idx) == 0 { return -1 } + return idx[0] +} +func (s *arrowDataFrameSchema) HasName(name string) bool { return len(s.schema.FieldIndices(name)) > 0 } +// Get from existing df.go (blue4209211/pq/df) returns df.SeriesSchema directly +func (s *arrowDataFrameSchema) Get(i int) df.SeriesSchema { + if i < 0 || i >= s.schema.NumFields() { panic("Get: index out of bounds for schema fields") } + field := s.schema.Field(i) + return df.SeriesSchema{Name: field.Name, Format: ArrowToDfFormat(field.Type), Nullable: field.Nullable} +} +func (s *arrowDataFrameSchema) Len() int { return s.schema.NumFields() } +func (s *arrowDataFrameSchema) Equals(other df.DataFrameSchema) bool { + if other == nil { return false } + otherArrowSchema, ok := other.(*arrowDataFrameSchema) + if !ok { + if s.Len() != other.Len() { return false } + for i := 0; i < s.Len(); i++ { + s1, s2 := s.Get(i), other.Get(i) + // Assuming df.Format has an Equals method or is comparable. + formatEquals := false + if s1.Format != nil && other.Get(i).Format != nil { + // This check is problematic if df.Format is an interface. + // Using reflect.DeepEqual for formats as a general approach if Equals method not on interface. + formatEquals = reflect.DeepEqual(s1.Format, s2.Format) + } else if s1.Format == nil && other.Get(i).Format == nil { + formatEquals = true + } + if s1.Name != s2.Name || !formatEquals || s1.Nullable != s2.Nullable { return false } + } + return true + } + return s.schema.Equal(otherArrowSchema.schema) +} +var _ df.DataFrameSchema = (*arrowDataFrameSchema)(nil) + + +// --- Helper Functions (package level) --- +// Exported ArrowToDfFormat for use in tests +func ArrowToDfFormat(dt arrow.DataType) df.Format { + switch dt.ID() { + case arrow.STRING, arrow.LARGE_STRING: return df.StringFormat + case arrow.BINARY, arrow.LARGE_BINARY: return df.StringFormat // Consider a distinct df.BinaryFormat if needed + case arrow.INT8, arrow.INT16, arrow.INT32, arrow.INT64: return df.IntegerFormat + case arrow.UINT8, arrow.UINT16, arrow.UINT32, arrow.UINT64: return df.IntegerFormat + case arrow.FLOAT32, arrow.FLOAT64: return df.DoubleFormat + case arrow.BOOL: return df.BoolFormat + case arrow.TIMESTAMP: return df.DateTimeFormat + case arrow.DATE32, arrow.DATE64: return df.DateTimeFormat + case arrow.NULL: return df.UnknownFormat + default: + panic(fmt.Sprintf("ArrowToDfFormat: unhandled Arrow data type %s", dt.Name())) + } +} +func dfFormatToArrowType(f df.Format) (arrow.DataType, error) { // Made error return consistent + if f == nil { return nil, fmt.Errorf("dfFormatToArrowType: df.Format cannot be nil") } + switch f { // Assuming f is comparable (not an interface with different underlying types for same logical format) + case df.StringFormat: return arrow.BinaryTypes.String, nil + case df.IntegerFormat: return arrow.PrimitiveTypes.Int64, nil + case df.DoubleFormat: return arrow.PrimitiveTypes.Float64, nil + case df.BoolFormat: return arrow.PrimitiveTypes.Boolean, nil + case df.DateTimeFormat: return arrow.TimestampTypes.Timestamp_ns, nil + case df.UnknownFormat: return nil, fmt.Errorf("cannot convert df.UnknownFormat to Arrow type") + default: + // This path means f is a df.Format instance not matching known singletons. + // It might be a custom format or an issue with how formats are defined/compared. + return nil, fmt.Errorf("dfFormatToArrowType: unhandled df.Format name %s, type %v", f.Name(), f.Type()) + } +} + +func dfValueToArrowScalar(val df.Value, targetType arrow.DataType, mem memory.Allocator) (scalar.Scalar, error) { + if val == nil { return nil, fmt.Errorf("input df.Value is nil interface") } + + var effectiveTargetType arrow.DataType = targetType + if effectiveTargetType == nil { + var err error + effectiveTargetType, err = dfFormatToArrowType(val.Schema()) + if err != nil { + return nil, fmt.Errorf("dfValueToArrowScalar: cannot determine target Arrow type for df.Value schema %v: %w", val.Schema(), err) + } + } + + if val.IsNil() { + return scalar.NewNullScalar(effectiveTargetType), nil + } + + if av, ok := val.(*arrowValue); ok { + if arrow.TypeEqual(av.val.DataType(), effectiveTargetType) { return av.val, nil } + ctx := context.Background(); if mem != nil { ctx = compute.WithAllocator(ctx, mem) } + castedScalar, err := scalar.Cast(ctx, av.val, effectiveTargetType) + if err != nil { return nil, fmt.Errorf("failed to cast scalar from %s to %s: %w", av.val.DataType(), effectiveTargetType, err) } + return castedScalar, nil + } + + // Fallback for generic df.Value + // This is less type-safe and relies on Get() returning types that match the df.Format's promise. + switch val.Schema() { // Use val.Schema() to guide conversion from generic df.Value + case df.IntegerFormat: + return scalar.NewInt64Scalar(val.GetAsInt()), nil + case df.DoubleFormat: + return scalar.NewFloat64Scalar(val.GetAsDouble()), nil + case df.StringFormat: + return scalar.NewStringScalar(val.GetAsString()), nil + case df.BoolFormat: + return scalar.NewBooleanScalar(val.GetAsBool()), nil + case df.DateTimeFormat: + if tsType, ok := effectiveTargetType.(*arrow.TimestampType); ok { + return scalar.NewTimestampFromTime(val.GetAsDatetime(), tsType), nil + } + return nil, fmt.Errorf("targetType for DateTimeFormat must be TimestampType, got %s", effectiveTargetType.Name()) + default: + return nil, fmt.Errorf("unsupported conversion from df.Value (format %s, Go type %T) to Arrow type %s", val.Schema().Name(), val.Get(), effectiveTargetType.Name()) + } +} + +func appendScalarToBuilder(b array.Builder, s scalar.Scalar, targetType arrow.DataType) error { + if s == nil { return fmt.Errorf("appendScalarToBuilder: input scalar is nil pointer")} + if !s.IsValid() { b.AppendNull(); return nil } + + var scalarToAppend scalar.Scalar = s + var castedScalarReleaser memory.Releasable + + if !arrow.TypeEqual(s.DataType(), targetType) { + ctx := context.Background() + casted, err := scalar.Cast(ctx, s, targetType) + if err != nil { + return fmt.Errorf("casting scalar from %s to %s failed: %w", s.DataType(), targetType, err) + } + scalarToAppend = casted + if releasable, ok := casted.(memory.Releasable); ok { + castedScalarReleaser = releasable + } + } + if castedScalarReleaser != nil { + defer castedScalarReleaser.Release() + } + + switch tb := b.(type) { + case *builder.Int64Builder: + v, ok := scalarToAppend.(*scalar.Int64); if !ok { return fmt.Errorf("expected Int64 scalar for Int64Builder, got %T (value: %s)", scalarToAppend, scalarToAppend) }; tb.Append(v.Value) + case *builder.Float64Builder: + v, ok := scalarToAppend.(*scalar.Float64); if !ok { return fmt.Errorf("expected Float64 scalar, got %T", scalarToAppend) }; tb.Append(v.Value) + case *builder.StringBuilder: + v, ok := scalarToAppend.(scalar.StringScalar); if !ok { return fmt.Errorf("expected StringScalar, got %T", scalarToAppend) }; tb.Append(v.String()) + case *builder.BooleanBuilder: + v, ok := scalarToAppend.(*scalar.Boolean); if !ok { return fmt.Errorf("expected Boolean scalar, got %T", scalarToAppend) }; tb.Append(v.Value) + case *builder.TimestampBuilder: + v, ok := scalarToAppend.(*scalar.Timestamp); if !ok { return fmt.Errorf("expected Timestamp scalar, got %T", scalarToAppend) }; tb.Append(v.Value) + case *builder.Date32Builder: + v, ok := scalarToAppend.(*scalar.Date32); if !ok { return fmt.Errorf("expected Date32 scalar, got %T", scalarToAppend) }; tb.Append(v.Value) + case *builder.Date64Builder: + v, ok := scalarToAppend.(*scalar.Date64); if !ok { return fmt.Errorf("expected Date64 scalar, got %T", scalarToAppend) }; tb.Append(v.Value) + default: + return fmt.Errorf("unsupported builder type in appendScalarToBuilder: %T for scalar type %s", b, scalarToAppend.DataType().Name()) + } + return nil +} + +// Helper functions to check df.Format type. +func isConcreteFormatType(format df.Format, targetKnownType df.Format) bool { + if format == nil { return false } + return format == targetKnownType +} +// These assume df package defines these as comparable values (e.g. var StringFormat = &formatImpl{...}) +func isIntegerFormat(f df.Format) bool { return isConcreteFormatType(f, df.IntegerFormat) } +func isFloatFormat(f df.Format) bool { return isConcreteFormatType(f, df.DoubleFormat) } +func isStringFormat(f df.Format) bool { return isConcreteFormatType(f, df.StringFormat) } +func isBoolFormat(f df.Format) bool { return isConcreteFormatType(f, df.BoolFormat) } +func isTimeFormat(f df.Format) bool { return isConcreteFormatType(f, df.DateTimeFormat) } +func isDateFormat(f df.Format) bool { + // Assuming DateTimeFormat might be used for dates if no specific DateFormat exists in df, + // or if df.DateFormat is distinct. This depends on df package definitions. + // If df.DateFormat exists and is distinct from df.DateTimeFormat: + // return isConcreteFormatType(f, df.DateFormat) || isConcreteFormatType(f, df.DateTimeFormat) // If DateTime can hold dates + return isConcreteFormatType(f, df.DateTimeFormat) // Or a specific df.DateFormat +} + + +// NewNilArrowRow creates a df.Row where all values are nil, matching the given schema. +// Exported for direct use in tests or other packages if necessary. +func NewNilArrowRow(schema *arrowDataFrameSchema, mem memory.Allocator) df.Row { + if schema == nil || schema.schema == nil { + panic("NewNilArrowRow: input schema or its internal arrow.Schema cannot be nil") + } + // mem is not directly used by NewNullScalar or newArrowRow currently. + if schema.schema.NumFields() == 0 { + return newArrowRow(schema, []scalar.Scalar{}) + } + + nilScalars := make([]scalar.Scalar, schema.schema.NumFields()) + for i, field := range schema.schema.Fields() { + nilScalars[i] = scalar.NewNullScalar(field.Type) + } + return newArrowRow(schema, nilScalars) +} +``` +And `df/arrow/types_test.go` with the new test: +```go +//go:build arrow + +package arrow_test + +import ( + "fmt" + "reflect" + "sort" + "strconv" + "testing" + "time" + + "git.querycap.com/practice/df" + arrowimpl "git.querycap.com/practice/df/arrow" + "github.com/apache/arrow/go/v14/arrow" + "github.com/apache/arrow/go/v14/arrow/array" + "github.com/apache/arrow/go/v14/arrow/memory" + "github.com/apache/arrow/go/v14/arrow/scalar" + "github.com/stretchr/testify/assert" +) + +// Helper to get a common test schema (assumed to exist from prior tests) +func getTestArrowSchemaForNilRowTest() *arrow.Schema { + return arrow.NewSchema( + []arrow.Field{ + {Name: "col_str", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "col_int", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "col_float", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, + {Name: "col_bool", Type: arrow.PrimitiveTypes.Boolean, Nullable: true}, + {Name: "col_time_ns", Type: arrow.TimestampTypes.Timestamp_ns, Nullable: true}, + {Name: "col_date32", Type: arrow.FixedWidthTypes.Date32, Nullable: true}, + }, + nil, + ) +} + + +func TestNewNilArrowRow(t *testing.T) { + mem := memory.NewGoAllocator() + + t.Run("SchemaWithMultipleFields", func(t *testing.T) { + arrowSchema := getTestArrowSchemaForNilRowTest() + // Cast to internal type *arrowimpl.ArrowDataFrameSchema for NewNilArrowRow + // NewArrowDataFrameSchema is the public constructor. + dfSchema := arrowimpl.NewArrowDataFrameSchema(arrowSchema).(*arrowimpl.ArrowDataFrameSchema) + + // Call the (now exported for test) NewNilArrowRow + nilRow := arrowimpl.NewNilArrowRow(dfSchema, mem) + + assert.NotNil(t, nilRow) + assert.Equal(t, arrowSchema.NumFields(), nilRow.Len(), "Row length should match field count") + + assert.True(t, nilRow.Schema().Equals(dfSchema), "Row schema should match input dfSchema") + + for i := 0; i < nilRow.Len(); i++ { + originalField := arrowSchema.Field(i) + // df.Row.IsNil does not return error + assert.True(t, nilRow.IsNil(i), fmt.Sprintf("Value at index %d (%s) should be nil", i, originalField.Name)) + + // df.Row.Get does not return error + val := nilRow.Get(i) + assert.NotNil(t, val, "df.Value should not be nil interface") + assert.True(t, val.IsNil(), fmt.Sprintf("df.Value at index %d (%s) should be nil", i, originalField.Name)) + + expectedDfFormat := arrowimpl.ArrowToDfFormat(originalField.Type) + valSchemaFormat := val.Schema() // This is df.Format as per problem's df.Value interface + + // Check Name and Type of df.Format + // Need to ensure df.Format has Name() and Type() methods as per definition + // For this test, we assume df.Format is comparable or has an Equals method. + // If df.Format's Name() or Type() methods are not part of the actual df.Format interface, + // then this part of the test needs adjustment. + // For now, comparing directly assuming they are comparable or singletons from df package. + assert.Equal(t, expectedDfFormat, valSchemaFormat, fmt.Sprintf("Schema format for col %d (%s)", i, originalField.Name)) + + + // Also check the underlying scalar type in the arrowValue + if av, ok := val.(*arrowimpl.ArrowValue); ok { + assert.True(t, av.IsNilInternalScalar(), fmt.Sprintf("Internal scalar for col %d (%s) should be nil/invalid", i, originalField.Name)) + assert.True(t, arrow.TypeEqual(originalField.Type, av.InternalScalarType()), fmt.Sprintf("Internal scalar type for col %d (%s) should match field type", i, originalField.Name)) + } else { + t.Errorf("nilRow.Get(%d) did not return an *arrowimpl.ArrowValue as expected by this test's internals", i) + } + } + }) + + t.Run("SchemaWithZeroFields", func(t *testing.T) { + arrowSchema0 := arrow.NewSchema([]arrow.Field{}, nil) + dfSchema0 := arrowimpl.NewArrowDataFrameSchema(arrowSchema0).(*arrowimpl.ArrowDataFrameSchema) + nilRow0 := arrowimpl.NewNilArrowRow(dfSchema0, mem) + assert.NotNil(t, nilRow0) + assert.Equal(t, 0, nilRow0.Len(), "Row length for 0-field schema should be 0") + assert.True(t, nilRow0.Schema().Equals(dfSchema0)) + }) + + t.Run("PanicOnNilSchema", func(t *testing.T) { + var nilDfSchema *arrowimpl.ArrowDataFrameSchema = nil + assert.PanicsWithValue(t, "NewNilArrowRow: input schema or its internal arrow.Schema cannot be nil", func() { + arrowimpl.NewNilArrowRow(nilDfSchema, mem) + }) + }) +} + +// Placeholder for other tests that might exist in this file +func TestArrowValue(t *testing.T) { + // Example of how other tests might look + t.Run("StringValue", func(t *testing.T) { + s := scalar.NewStringScalar("hello") + v := arrowimpl.NewArrowValue(s, df.StringFormat) // Assuming df.StringFormat is defined + assert.False(t, v.IsNil()) + assert.Equal(t, "hello", v.Get()) + assert.Equal(t, "hello", v.GetAsString()) + assert.Equal(t, df.StringFormat, v.Schema()) + }) +} + +func TestArrowDataFrameSchema(t *testing.T) { + // Example + t.Run("EmptySchema", func(t *testing.T) { + as := arrow.NewSchema(nil,nil) + dfs := arrowimpl.NewArrowDataFrameSchema(as) + assert.Equal(t, 0, dfs.Len()) + }) +} +``` +df/arrow/types_test.go +<<<<<<< SEARCH +//go:build arrow + +package arrow_test + +import ( + "reflect" + "testing" + "time" + + "github.com/apache/arrow/go/v14/arrow" + "github.com/apache/arrow/go/v14/arrow/array" + "github.com/apache/arrow/go/v14/arrow/memory" + "github.com/apache/arrow/go/v14/arrow/scalar" + "git.querycap.com/practice/df" // Ensure this is the correct import path + "github.com/stretchr/testify/assert" + + arrowimpl "git.querycap.com/practice/df/arrow" // Package to test (internal call) +) + +// --- (Existing tests for arrowValue, arrowDataFrameSchema, arrowRow must be here) --- +// --- (Helper getTestArrowSchema() from previous tests is useful here) --- +/* +func getTestArrowSchema() *arrow.Schema { // From previous tests + return arrow.NewSchema( + []arrow.Field{ + {Name: "col_str", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "col_int", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "col_float", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, + {Name: "col_bool", Type: arrow.PrimitiveTypes.Boolean, Nullable: true}, + {Name: "col_time", Type: arrow.TimestampTypes.Timestamp_ns, Nullable: true}, + }, + nil, + ) +} +*/ + + +func TestNewNilArrowRow(t *testing.T) { + mem := memory.NewGoAllocator() // Allocator might not be strictly needed by NewNullScalar but good practice if helpers use it. + + // Case 1: Schema with multiple fields of different types + arrowSchema1 := arrow.NewSchema( + []arrow.Field{ + {Name: "c1_str", Type: arrow.BinaryTypes.String}, + {Name: "c2_int", Type: arrow.PrimitiveTypes.Int64}, + {Name: "c3_bool", Type: arrow.PrimitiveTypes.Boolean}, + }, nil, + ) + dfSchema1 := arrowimpl.NewArrowDataFrameSchema(arrowSchema1).(*arrowimpl.ArrowDataFrameSchema) + + nilRow1 := arrowimpl.CallNewNilArrowRow(dfSchema1, mem) // Using exported wrapper if newNilArrowRow is not exported + + assert.NotNil(t, nilRow1) + assert.Equal(t, 3, nilRow1.Len(), "Row length should match field count") + assert.True(t, nilRow1.Schema().Equals(dfSchema1), "Row schema should match input") + + for i := 0; i < nilRow1.Len(); i++ { + assert.True(t, nilRow1.IsNil(i), fmt.Sprintf("Value at index %d should be nil", i)) + val := nilRow1.Get(i) + assert.True(t, val.IsNil(), fmt.Sprintf("df.Value at index %d should be nil", i)) + // Check that the df.Value schema format matches the original field type + expectedDfFormat := arrowimpl.ArrowToDfFormat(arrowSchema1.Field(i).Type) + assert.Equal(t, expectedDfFormat.Name(), val.Schema().Name(), fmt.Sprintf("Schema name for col %d", i)) + assert.Equal(t, expectedDfFormat.Type(), val.Schema().Type(), fmt.Sprintf("Schema type for col %d", i)) + } + + // Case 2: Schema with zero fields + arrowSchema0 := arrow.NewSchema([]arrow.Field{}, nil) + dfSchema0 := arrowimpl.NewArrowDataFrameSchema(arrowSchema0).(*arrowimpl.ArrowDataFrameSchema) + nilRow0 := arrowimpl.CallNewNilArrowRow(dfSchema0, mem) + assert.NotNil(t, nilRow0) + assert.Equal(t, 0, nilRow0.Len(), "Row length for 0-field schema should be 0") + assert.True(t, nilRow0.Schema().Equals(dfSchema0)) + + + // Case 3: Panic if input schema is nil or its internal arrow.Schema is nil + var nilDfSchema *arrowimpl.ArrowDataFrameSchema = nil + assert.PanicsWithValue(t, "newNilArrowRow: input schema or its internal arrow.Schema cannot be nil", func() { + arrowimpl.CallNewNilArrowRow(nilDfSchema, mem) + }) + + // Create a dfSchema with a nil internal arrow.Schema + // This state should ideally not occur if constructors are used correctly. + // Forcing it for test: + dfSchemaWithNilInternal := arrowimpl.NewArrowDataFrameSchema(nil).(*arrowimpl.ArrowDataFrameSchema) + // The above NewArrowDataFrameSchema(nil) actually creates a valid empty schema. + // So the panic "internal arrow.Schema cannot be nil" inside newNilArrowRow might not be reachable + // if dfSchema itself is not nil. The first check `schema == nil` covers nilDfSchema. + // The check `schema.schema == nil` inside newNilArrowRow: + // If NewArrowDataFrameSchema(nil) creates schema.schema as non-nil empty schema, then this is fine. +} + +// --- (Rest of types_test.go content) --- +======= +//go:build arrow + +package arrow_test + +import ( + "fmt" + "reflect" + "sort" + "strconv" + "testing" + "time" + + "git.querycap.com/practice/df" + arrowimpl "git.querycap.com/practice/df/arrow" + "github.com/apache/arrow/go/v14/arrow" + "github.com/apache/arrow/go/v14/arrow/array" + "github.com/apache/arrow/go/v14/arrow/memory" + "github.com/apache/arrow/go/v14/arrow/scalar" + "github.com/stretchr/testify/assert" +) + +// Helper to get a common test schema (assumed to exist from prior tests) +func getTestArrowSchemaForNilRowTest() *arrow.Schema { + return arrow.NewSchema( + []arrow.Field{ + {Name: "col_str", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "col_int", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "col_float", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, + {Name: "col_bool", Type: arrow.PrimitiveTypes.Boolean, Nullable: true}, + {Name: "col_time_ns", Type: arrow.TimestampTypes.Timestamp_ns, Nullable: true}, + {Name: "col_date32", Type: arrow.FixedWidthTypes.Date32, Nullable: true}, + }, + nil, + ) +} + + +func TestNewNilArrowRow(t *testing.T) { + mem := memory.NewGoAllocator() + + t.Run("SchemaWithMultipleFields", func(t *testing.T) { + arrowSchema := getTestArrowSchemaForNilRowTest() + // Cast to internal type *arrowimpl.ArrowDataFrameSchema for NewNilArrowRow + // NewArrowDataFrameSchema is the public constructor. + dfSchema := arrowimpl.NewArrowDataFrameSchema(arrowSchema).(*arrowimpl.ArrowDataFrameSchema) + + // Call the (now exported for test) NewNilArrowRow + nilRow := arrowimpl.NewNilArrowRow(dfSchema, mem) + + assert.NotNil(t, nilRow) + assert.Equal(t, arrowSchema.NumFields(), nilRow.Len(), "Row length should match field count") + + assert.True(t, nilRow.Schema().Equals(dfSchema), "Row schema should match input dfSchema") + + for i := 0; i < nilRow.Len(); i++ { + originalField := arrowSchema.Field(i) + // df.Row.IsNil does not return error + assert.True(t, nilRow.IsNil(i), fmt.Sprintf("Value at index %d (%s) should be nil", i, originalField.Name)) + + // df.Row.Get does not return error + val := nilRow.Get(i) + assert.NotNil(t, val, "df.Value should not be nil interface") + assert.True(t, val.IsNil(), fmt.Sprintf("df.Value at index %d (%s) should be nil", i, originalField.Name)) + + expectedDfFormat := arrowimpl.ArrowToDfFormat(originalField.Type) + valSchemaFormat := val.Schema() // This is df.Format as per problem's df.Value interface + + assert.Equal(t, expectedDfFormat, valSchemaFormat, fmt.Sprintf("Schema format for col %d (%s)", i, originalField.Name)) + + + // Also check the underlying scalar type in the arrowValue + if av, ok := val.(*arrowimpl.ArrowValue); ok { + assert.True(t, av.IsNilInternalScalar(), fmt.Sprintf("Internal scalar for col %d (%s) should be nil/invalid", i, originalField.Name)) + assert.True(t, arrow.TypeEqual(originalField.Type, av.InternalScalarType()), fmt.Sprintf("Internal scalar type for col %d (%s) should match field type", i, originalField.Name)) + } else { + t.Errorf("nilRow.Get(%d) did not return an *arrowimpl.ArrowValue as expected by this test's internals", i) + } + } + }) + + t.Run("SchemaWithZeroFields", func(t *testing.T) { + arrowSchema0 := arrow.NewSchema([]arrow.Field{}, nil) + dfSchema0 := arrowimpl.NewArrowDataFrameSchema(arrowSchema0).(*arrowimpl.ArrowDataFrameSchema) + nilRow0 := arrowimpl.NewNilArrowRow(dfSchema0, mem) + assert.NotNil(t, nilRow0) + assert.Equal(t, 0, nilRow0.Len(), "Row length for 0-field schema should be 0") + assert.True(t, nilRow0.Schema().Equals(dfSchema0)) + }) + + t.Run("PanicOnNilSchema", func(t *testing.T) { + var nilDfSchema *arrowimpl.ArrowDataFrameSchema = nil + assert.PanicsWithValue(t, "NewNilArrowRow: input schema or its internal arrow.Schema cannot be nil", func() { + arrowimpl.NewNilArrowRow(nilDfSchema, mem) + }) + }) +} + +// Placeholder for other tests that might exist in this file +func TestArrowValue(t *testing.T) { + // Example of how other tests might look + t.Run("StringValue", func(t *testing.T) { + s := scalar.NewStringScalar("hello") + // Assuming df.StringFormat is defined and accessible + // For testing, ensure df.StringFormat is part of the df package or defined for tests. + // If df.StringFormat is an interface, this comparison might not be direct. + // Let's use a format that's known to be df.StringFormat for this test. + + // Get the df.Format by converting from known Arrow type + strFormat := arrowimpl.ArrowToDfFormat(arrow.BinaryTypes.String) + + v := arrowimpl.NewArrowValue(s, strFormat) + assert.False(t, v.IsNil()) + assert.Equal(t, "hello", v.Get()) + assert.Equal(t, "hello", v.GetAsString()) + assert.Equal(t, strFormat, v.Schema()) + }) +} + +func TestArrowDataFrameSchema(t *testing.T) { + // Example + t.Run("EmptySchema", func(t *testing.T) { + as := arrow.NewSchema(nil,nil) // Creates an empty schema + dfs := arrowimpl.NewArrowDataFrameSchema(as) + assert.Equal(t, 0, dfs.Len()) + }) +} +>>>>>>> REPLACE diff --git a/df/arrow/types_test.go b/df/arrow/types_test.go new file mode 100644 index 0000000..1afc4d3 --- /dev/null +++ b/df/arrow/types_test.go @@ -0,0 +1,410 @@ +//go:build arrow + +package arrow_test + +import ( + "reflect" + "testing" + "time" + + "github.com/apache/arrow/go/v14/arrow" + "github.com/apache/arrow/go/v14/arrow/array" + "github.com/apache/arrow/go/v14/arrow/memory" + "github.com/apache/arrow/go/v14/arrow/scalar" + "github.com/blue4209211/pq/df" + "github.com/stretchr/testify/assert" + + arrowimpl "github.com/blue4209211/pq/df/arrow" // Import the implementation package +) + +func TestArrowValue_String(t *testing.T) { + s := scalar.NewStringScalar("hello") + f := df.StringFormat + val := arrowimpl.NewArrowValue(s, f) + + assert.Equal(t, f, val.Schema()) + assert.Equal(t, "hello", val.Get()) + assert.Equal(t, "hello", val.GetAsString()) + assert.False(t, val.IsNil()) + + sNil := scalar.NewNullScalar(arrow.BinaryTypes.String) + valNil := arrowimpl.NewArrowValue(sNil, f) + assert.True(t, valNil.IsNil()) + assert.Equal(t, "", valNil.GetAsString()) // Behavior for nil GetAsString + + // Equals + s2 := scalar.NewStringScalar("hello") + val2 := arrowimpl.NewArrowValue(s2, f) + assert.True(t, val.Equals(val2)) + + s3 := scalar.NewStringScalar("world") + val3 := arrowimpl.NewArrowValue(s3, f) + assert.False(t, val.Equals(val3)) + assert.False(t, val.Equals(nil)) + assert.False(t, valNil.Equals(val)) + assert.True(t, valNil.Equals(arrowimpl.NewArrowValue(scalar.NewNullScalar(arrow.BinaryTypes.String), f))) + + // Test Get for other types (should panic or handle error) + assert.Panics(t, func() { val.GetAsInt() }) +} + +func TestArrowValue_Int64(t *testing.T) { + s := scalar.NewInt64Scalar(123) + f := df.IntegerFormat + val := arrowimpl.NewArrowValue(s, f) + + assert.Equal(t, f, val.Schema()) + assert.Equal(t, int64(123), val.Get()) + assert.Equal(t, int64(123), val.GetAsInt()) + assert.Equal(t, "123", val.GetAsString()) + assert.False(t, val.IsNil()) + + // Equals + s2 := scalar.NewInt64Scalar(123) + val2 := arrowimpl.NewArrowValue(s2, f) + assert.True(t, val.Equals(val2)) + + s3 := scalar.NewInt64Scalar(456) + val3 := arrowimpl.NewArrowValue(s3, f) + assert.False(t, val.Equals(val3)) +} + +func TestArrowValue_Float64(t *testing.T) { + s := scalar.NewFloat64Scalar(123.456) + f := df.DoubleFormat + val := arrowimpl.NewArrowValue(s, f) + + assert.Equal(t, f, val.Schema()) + assert.Equal(t, 123.456, val.Get()) + assert.Equal(t, 123.456, val.GetAsDouble()) + assert.Equal(t, "123.456", val.GetAsString()) // Behavior of fmt.Sprintf might vary + assert.False(t, val.IsNil()) +} + +func TestArrowValue_Boolean(t *testing.T) { + sTrue := scalar.NewBooleanScalar(true) + f := df.BoolFormat + valTrue := arrowimpl.NewArrowValue(sTrue, f) + + assert.Equal(t, f, valTrue.Schema()) + assert.Equal(t, true, valTrue.Get()) + assert.Equal(t, true, valTrue.GetAsBool()) + assert.Equal(t, "true", valTrue.GetAsString()) + assert.False(t, valTrue.IsNil()) + + sFalse := scalar.NewBooleanScalar(false) + valFalse := arrowimpl.NewArrowValue(sFalse, f) + assert.Equal(t, false, valFalse.GetAsBool()) +} + +func TestArrowValue_Timestamp(t *testing.T) { + now := time.Now().Truncate(time.Nanosecond) + tsType := arrow.TimestampTypes.Timestamp_ns + s := scalar.NewTimestampScalar(arrow.Timestamp(now.UnixNano()), tsType) + f := df.DateTimeFormat + val := arrowimpl.NewArrowValue(s, f) + + assert.Equal(t, f, val.Schema()) + retrievedTime := val.Get().(time.Time) + assert.Equal(t, now.UnixNano(), retrievedTime.UnixNano()) + assert.True(t, now.Equal(val.GetAsDatetime())) + assert.False(t, val.IsNil()) +} + +func TestArrowValue_Equals_DifferentTypes(t *testing.T) { + sInt := scalar.NewInt64Scalar(10) + fInt := df.IntegerFormat + valInt := arrowimpl.NewArrowValue(sInt, fInt) + + sStr := scalar.NewStringScalar("10") + fStr := df.StringFormat + valStr := arrowimpl.NewArrowValue(sStr, fStr) + + assert.False(t, valInt.Equals(valStr), "Values of different underlying types should not be equal") + + mockOtherValue := &mockValue{data: int64(10), format: df.IntegerFormat} + assert.False(t, valInt.Equals(mockOtherValue), "arrowValue should not equal a different df.Value implementation by default") + +} + +type mockValue struct { + data any + format df.Format + isNil bool +} + +func (m *mockValue) Schema() df.Format { return m.format } +func (m *mockValue) Get() any { return m.data } +func (m *mockValue) GetAsString() string { v, _ := m.data.(string); return v } +func (m *mockValue) GetAsInt() int64 { v, _ := m.data.(int64); return v } +func (m *mockValue) GetAsDouble() float64 { v, _ := m.data.(float64); return v } +func (m *mockValue) GetAsBool() bool { v, _ := m.data.(bool); return v } +func (m *mockValue) GetAsDatetime() time.Time { v, _ := m.data.(time.Time); return v } +func (m *mockValue) IsNil() bool { return m.isNil } +func (m *mockValue) Equals(other df.Value) bool { + if other == nil { return m.isNil } + if m.IsNil() != other.IsNil() { return false } + if m.IsNil() && other.IsNil() { return true} + return reflect.DeepEqual(m.Get(), other.Get()) && m.Schema().Name() == other.Schema().Name() +} + +func getTestArrowSchema() *arrow.Schema { + return arrow.NewSchema( + []arrow.Field{ + {Name: "col_str", Type: arrow.BinaryTypes.String}, + {Name: "col_int", Type: arrow.PrimitiveTypes.Int64}, + {Name: "col_float", Type: arrow.PrimitiveTypes.Float64}, + {Name: "col_bool", Type: arrow.PrimitiveTypes.Boolean}, + {Name: "col_time", Type: arrow.TimestampTypes.Timestamp_ns}, + }, + nil, + ) +} + +func TestArrowDataFrameSchema_Basic(t *testing.T) { + arrowSchema := getTestArrowSchema() + dfSchema := arrowimpl.NewArrowDataFrameSchema(arrowSchema) + + assert.Equal(t, 5, dfSchema.Len()) + assert.Equal(t, []string{"col_str", "col_int", "col_float", "col_bool", "col_time"}, dfSchema.Names()) + seriesSchema0 := dfSchema.Get(0) + assert.Equal(t, "col_str", seriesSchema0.Name) + assert.Equal(t, df.StringFormat.Name(), seriesSchema0.Format.Name()) + seriesSchema1 := dfSchema.Get(1) + assert.Equal(t, "col_int", seriesSchema1.Name) + assert.Equal(t, df.IntegerFormat.Name(), seriesSchema1.Format.Name()) + seriesSchemaStr := dfSchema.GetByName("col_str") + assert.Equal(t, "col_str", seriesSchemaStr.Name) + assert.Equal(t, df.StringFormat.Name(), seriesSchemaStr.Format.Name()) + seriesSchemaFloat := dfSchema.GetByName("col_float") + assert.Equal(t, "col_float", seriesSchemaFloat.Name) + assert.Equal(t, df.DoubleFormat.Name(), seriesSchemaFloat.Format.Name()) + assert.Equal(t, -1, dfSchema.GetIndexByName("non_existent_col")) + assert.Equal(t, 0, dfSchema.GetIndexByName("col_str")) + assert.Equal(t, 2, dfSchema.GetIndexByName("col_float")) + assert.True(t, dfSchema.HasName("col_bool")) + assert.False(t, dfSchema.HasName("non_existent_col")) + allSeries := dfSchema.Series() + assert.Equal(t, 5, len(allSeries)) + assert.Equal(t, "col_time", allSeries[4].Name) + assert.Equal(t, df.DateTimeFormat.Name(), allSeries[4].Format.Name()) +} + +func TestArrowDataFrameSchema_Equals(t *testing.T) { + schema1 := arrowimpl.NewArrowDataFrameSchema(getTestArrowSchema()) + schema2 := arrowimpl.NewArrowDataFrameSchema(getTestArrowSchema()) + schemaDiffName := arrowimpl.NewArrowDataFrameSchema(arrow.NewSchema( + []arrow.Field{{Name: "col_str_diff", Type: arrow.BinaryTypes.String}, getTestArrowSchema().Field(1)}, nil)) + schemaDiffType := arrowimpl.NewArrowDataFrameSchema(arrow.NewSchema( + []arrow.Field{{Name: "col_str", Type: arrow.PrimitiveTypes.Int64}, getTestArrowSchema().Field(1)}, nil)) + schemaDiffLen := arrowimpl.NewArrowDataFrameSchema(arrow.NewSchema( + []arrow.Field{getTestArrowSchema().Field(0)}, nil)) + var nilArrowSchema *arrow.Schema = nil + schemaNilInternal := arrowimpl.NewArrowDataFrameSchema(nilArrowSchema) + var nilDfSchema df.DataFrameSchema = nil + + assert.True(t, schema1.Equals(schema2), "Identical schemas should be equal") + assert.False(t, schema1.Equals(schemaDiffName), "Schemas with different names should not be equal") + assert.False(t, schema1.Equals(schemaDiffType), "Schemas with different types should not be equal") + assert.False(t, schema1.Equals(schemaDiffLen), "Schemas with different lengths should not be equal") + assert.False(t, schema1.Equals(nilDfSchema), "Schema should not be equal to nil df.DataFrameSchema") + assert.False(t, schemaNilInternal.Equals(schema1), "Schema with nil internal arrow.Schema should not equal a valid one") + assert.True(t, schemaNilInternal.Equals(arrowimpl.NewArrowDataFrameSchema(nilArrowSchema)), "Two schemas with nil internal arrow.Schema should be equal") + + mockSchemaEq := &mockDataFrameSchema{ + series: []df.SeriesSchema{ + {Name: "col_str", Format: df.StringFormat}, {Name: "col_int", Format: df.IntegerFormat}, + {Name: "col_float", Format: df.DoubleFormat}, {Name: "col_bool", Format: df.BoolFormat}, + {Name: "col_time", Format: df.DateTimeFormat}, + }, + } + assert.True(t, schema1.Equals(mockSchemaEq), "arrowDataFrameSchema should be equal to a structurally equivalent mockDataFrameSchema") + + mockSchemaDiff := &mockDataFrameSchema{ + series: []df.SeriesSchema{{Name: "col_str", Format: df.StringFormat}, {Name: "col_int_diff", Format: df.IntegerFormat}}, + } + assert.False(t, schema1.Equals(mockSchemaDiff), "arrowDataFrameSchema should not be equal to a structurally different mockDataFrameSchema") +} + +type mockDataFrameSchema struct { + series []df.SeriesSchema +} +func (m *mockDataFrameSchema) Series() []df.SeriesSchema { return m.series } +func (m *mockDataFrameSchema) Names() []string { + names := make([]string, len(m.series)) + for i, s := range m.series { names[i] = s.Name } + return names +} +func (m *mockDataFrameSchema) GetByName(s string) df.SeriesSchema { + for _, ss := range m.series { if ss.Name == s { return ss } } + return df.SeriesSchema{} +} +func (m *mockDataFrameSchema) GetIndexByName(s string) int { + for i, ss := range m.series { if ss.Name == s { return i } } + return -1 +} +func (m *mockDataFrameSchema) HasName(s string) bool { + for _, ss := range m.series { if ss.Name == s { return true } } + return false +} +func (m *mockDataFrameSchema) Get(i int) df.SeriesSchema { return m.series[i] } +func (m *mockDataFrameSchema) Len() int { return len(m.series) } +func (m *mockDataFrameSchema) Equals(other df.DataFrameSchema) bool { + if other == nil || m.Len() != other.Len() { return false } + for i := 0; i < m.Len(); i++ { + s1 := m.Get(i) + s2 := other.Get(i) + if s1.Name != s2.Name || s1.Format.Name() != s2.Format.Name() || s1.Format.Type() != s2.Format.Type() { + return false + } + } + return true +} + +// --- Tests for arrowRow --- + +// Helper to create a sample record for testing arrowRow +func getTestRecord(mem memory.Allocator) arrow.Record { + schema := getTestArrowSchema() + b := array.NewRecordBuilder(mem, schema) + defer b.Release() + + b.Field(0).(*array.StringBuilder).AppendValues([]string{"hello", "world", "foo"}, nil) + b.Field(1).(*array.Int64Builder).AppendValues([]int64{1, 2, 0}, []bool{true, true, false}) // 0 is nil + b.Field(2).(*array.Float64Builder).AppendValues([]float64{1.1, 2.2, 3.3}, nil) + b.Field(3).(*array.BooleanBuilder).AppendValues([]bool{true, false, true}, nil) + now := time.Now().UnixNano() + b.Field(4).(*array.TimestampBuilder).AppendValues([]arrow.Timestamp{arrow.Timestamp(now), arrow.Timestamp(now + 1000), arrow.Timestamp(now + 2000)}, nil) + + return b.NewRecord() +} + + +func TestArrowRow_NewArrowRow(t *testing.T) { + schema := arrowimpl.NewArrowDataFrameSchema(getTestArrowSchema()).(*arrowimpl.ArrowDataFrameSchema) + vals := []scalar.Scalar{ + scalar.NewStringScalar("alpha"), + scalar.NewInt64Scalar(100), + scalar.NewFloat64Scalar(99.9), + scalar.NewBooleanScalar(true), + scalar.NewTimestampScalar(arrow.Timestamp(time.Now().UnixNano()), arrow.TimestampTypes.Timestamp_ns), + } + row := arrowimpl.NewArrowRow(schema, vals) + + assert.Equal(t, schema, row.Schema()) + assert.Equal(t, 5, row.Len()) + assert.Equal(t, "alpha", row.Get(0).GetAsString()) + assert.Equal(t, int64(100), row.Get(1).GetAsInt()) + assert.False(t, row.IsAnyNil()) +} + +func TestArrowRow_NewArrowRowFromRecord(t *testing.T) { + mem := memory.NewGoAllocator() + record := getTestRecord(mem) + defer record.Release() + + schema := arrowimpl.NewArrowDataFrameSchema(record.Schema()).(*arrowimpl.ArrowDataFrameSchema) + + // Test row 0 + row0, err0 := arrowimpl.NewArrowRowFromRecord(schema, record, 0) + assert.NoError(t, err0) + assert.Equal(t, schema, row0.Schema()) + assert.Equal(t, 5, row0.Len()) + assert.Equal(t, "hello", row0.Get(0).GetAsString()) + assert.Equal(t, "hello", row0.GetByName("col_str").GetAsString()) + assert.Equal(t, int64(1), row0.GetAsInt(1)) + assert.Equal(t, 1.1, row0.GetAsDouble(2)) + assert.True(t, row0.GetAsBool(3)) + assert.NotZero(t, row0.GetAsDatetime(4)) + assert.False(t, row0.IsNil(0)) + assert.False(t, row0.IsAnyNil()) + + // Test row 1 (with a nil value) + row1, err1 := arrowimpl.NewArrowRowFromRecord(schema, record, 1) + assert.NoError(t, err1) + assert.Equal(t, "world", row1.Get(0).GetAsString()) + assert.True(t, row1.IsNil(1), "col_int at index 1 should be nil") // Index 1 of col_int is nil + assert.True(t, row1.IsAnyNil()) + assert.Panics(t, func() { row1.GetAsInt(1) }, "GetAsInt on a nil value should panic") + + + // Test GetRaw + assert.Equal(t, "hello", row0.GetRaw(0)) + assert.Equal(t, int64(1), row0.GetRaw(1)) + + + // Test GetMap + rowMap := row0.GetMap() + assert.Equal(t, 5, len(rowMap)) + assert.Equal(t, "hello", rowMap["col_str"].GetAsString()) + assert.Equal(t, int64(1), rowMap["col_int"].GetAsInt()) + + // Test out of bounds + _, errBounds := arrowimpl.NewArrowRowFromRecord(schema, record, 10) + assert.Error(t, errBounds) + + assert.Panics(t, func() { row0.Get(10) }) + assert.Panics(t, func() { row0.GetByName("non_existent") }) +} + + +func TestArrowRow_Copy(t *testing.T) { + schema := arrowimpl.NewArrowDataFrameSchema(getTestArrowSchema()).(*arrowimpl.ArrowDataFrameSchema) + vals := []scalar.Scalar{scalar.NewStringScalar("copy_me"), scalar.NewInt64Scalar(55)} + + // Adjust schema to match vals + simpleArrowSchema := arrow.NewSchema([]arrow.Field{getTestArrowSchema().Field(0), getTestArrowSchema().Field(1)}, nil) + simpleDfSchema := arrowimpl.NewArrowDataFrameSchema(simpleArrowSchema).(*arrowimpl.ArrowDataFrameSchema) + + row := arrowimpl.NewArrowRow(simpleDfSchema, vals) + copiedRow := row.Copy() + + assert.True(t, row.Schema().Equals(copiedRow.Schema())) + assert.Equal(t, row.Len(), copiedRow.Len()) + assert.True(t, row.Get(0).Equals(copiedRow.Get(0))) + assert.True(t, row.Get(1).Equals(copiedRow.Get(1))) + + // Ensure it's a shallow copy of scalars (scalars are immutable-like) + // but the slice itself is new + originalVal0 := row.Get(0).(*arrowimpl.ArrowValue) + copiedVal0 := copiedRow.Get(0).(*arrowimpl.ArrowValue) + + // This checks if the underlying scalar.Scalar is the same instance. + // For simple scalars, this might be true due to how they are created/interned. + // The important part is that changes to one row's structure (if possible) wouldn't affect the other. + // Since our `arrowRow.values` is a slice of `scalar.Scalar`, `copy()` on the slice creates a new slice. + // And `scalar.Scalar` itself is an interface, the concrete types are typically pointers to structs + // that are value-based or immutable in nature. + + if len(vals) > 0 && originalVal0 != nil && copiedVal0 != nil { + // If Get returns an arrowValue, we can compare its internal scalar + // For simple scalars, they might point to the same underlying scalar instance if scalar.Copy() isn't used by MakeScalar or if scalars are interned. + // However, the slice `values` in `arrowRow` is copied. + } +} + +func TestArrowRow_Select(t *testing.T) { + mem := memory.NewGoAllocator() + record := getTestRecord(mem) + defer record.Release() + schema := arrowimpl.NewArrowDataFrameSchema(record.Schema()).(*arrowimpl.ArrowDataFrameSchema) + row, _ := arrowimpl.NewArrowRowFromRecord(schema, record, 0) + + // Select "col_str", "col_float" (indices 0, 2) + selectedRow := row.Select(0, 2) + assert.Equal(t, 2, selectedRow.Len()) + assert.Equal(t, "col_str", selectedRow.Schema().Get(0).Name) + assert.Equal(t, "col_float", selectedRow.Schema().Get(1).Name) + assert.Equal(t, "hello", selectedRow.Get(0).GetAsString()) + assert.Equal(t, 1.1, selectedRow.Get(1).GetAsDouble()) + + // Select in different order: "col_int", "col_str" (indices 1, 0) + selectedRowReordered := row.Select(1, 0) + assert.Equal(t, 2, selectedRowReordered.Len()) + assert.Equal(t, "col_int", selectedRowReordered.Schema().Get(0).Name) + assert.Equal(t, "col_str", selectedRowReordered.Schema().Get(1).Name) + assert.Equal(t, int64(1), selectedRowReordered.Get(0).GetAsInt()) + assert.Equal(t, "hello", selectedRowReordered.Get(1).GetAsString()) + + assert.Panics(t, func() { row.Select(0, 10) }, "Select with out-of-bounds index should panic") +} diff --git a/df/dataframe.go b/df/dataframe.go index 95842b0..679d8a2 100644 --- a/df/dataframe.go +++ b/df/dataframe.go @@ -104,6 +104,14 @@ type GroupedDataFrame interface { Map(f func(Row, DataFrame) DataFrame) GroupedDataFrame Where(f func(Row, DataFrame) bool) GroupedDataFrame Len() int64 + Agg(configs ...AggregationConfig) DataFrame +} + +// AggregationConfig defines an aggregation function, input column, and output column name. +type AggregationConfig struct { + Func string + InputCol string + OutputColName string } // Series Type for Storing column data of Dataframe diff --git a/df/fns/series/groups_test.go b/df/fns/series/groups_test.go index d46c2b7..ae0f2a2 100644 --- a/df/fns/series/groups_test.go +++ b/df/fns/series/groups_test.go @@ -46,15 +46,9 @@ func TestMedian(t *testing.T) { assert.Equal(t, float64(3), Median(s1).GetAsDouble()) } -func TestDescribe(t *testing.T) { -} - func TestCountDistinctValues(t *testing.T) { s1 := inmemory.NewIntSeriesVarArg(1, 2, 3, 4, 1, 1) s2 := CountDistinctValues(s1) assert.Equal(t, 4, len(s2)) assert.Equal(t, int64(3), s2["1"]) } - -func TestCovariance(t *testing.T) { -} diff --git a/df/inmemory/df.go b/df/inmemory/df.go index 976a9d3..15541ae 100644 --- a/df/inmemory/df.go +++ b/df/inmemory/df.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( @@ -201,8 +202,104 @@ func (t *inmemoryDataFrame) RenameSeriesByName(col string, name string, inplace return t.RenameSeries(index, name, inplace) } -func (t *inmemoryDataFrame) Select(index ...df.Expr) (d df.DataFrame) { - return d +func (t *inmemoryDataFrame) Select(expressions ...df.Expr) df.DataFrame { + if len(expressions) == 0 { + // Return a new DataFrame with the same number of rows but no columns + emptySchema := NewInMemorySchema(t.name, []df.SeriesSchema{}) + emptyRows := make([]df.Row, t.Len()) + for i := range emptyRows { + emptyRows[i] = NewInMemoryRow(emptySchema, []df.Value{}) + } + return NewDataframeFromRowAndName(t.name+"_select_empty", emptySchema, &emptyRows) + } + + var newSchemaSeries []df.SeriesSchema + newRowsData := make([]df.Row, 0, t.Len()) + var finalSchema df.DataFrameSchema + + for i, inputRow := range t.data { + outputValues := make([]df.Value, 0, len(expressions)) + for _, expr := range expressions { + var val df.Value + outputName := expr.Name() // Use expr.Name() for the output column name + + switch e := expr.(type) { + case df.ColNameExpr: + colName := e.Col() + val = inputRow.GetByName(colName) + if i == 0 { // First row, determine schema + // Ensure outputName is used for the schema + if outputName == "" { // Should ideally be set by Alias or Col itself + outputName = colName + } + newSchemaSeries = append(newSchemaSeries, df.SeriesSchema{Name: outputName, Format: val.Schema()}) + } + case df.LiteralExpr: + val = e.Const() + if i == 0 { // First row, determine schema + if outputName == "" { // Literals might not have a pre-defined name unless aliased + // Create a generic name or use a convention like "literal_N" + // For now, let's try to use the string representation of the literal, + Daunting if it's long. + // A better approach would be to require aliases for literals if a specific name is needed. + // For simplicity, if no alias, could panic or use a default. + // Let's assume expr.Name() handles aliasing correctly for literals too. + // If expr.Name() is empty for a literal, it implies it wasn't aliased. + // The problem statement implies expr.Name() should be used. + if outputName == "" { + panic(fmt.Sprintf("literal expression %v must have an alias via Name()", e.Const().GetAsString())) + } + } + newSchemaSeries = append(newSchemaSeries, df.SeriesSchema{Name: outputName, Format: val.Schema()}) + } + default: + // Placeholder for more complex expressions + panic(fmt.Sprintf("unsupported expression type: %T", expr)) + } + outputValues = append(outputValues, val) + } + + if i == 0 { + finalSchema = NewInMemorySchema(t.name+"_select", newSchemaSeries) + } + newRowsData = append(newRowsData, NewInMemoryRow(finalSchema, outputValues)) + } + + // Handle case where t.data is empty, newSchemaSeries would be empty too + if len(t.data) == 0 { + // Construct schema based on expressions, assuming types can be inferred + // without data. This is tricky for ColNameExpr without data. + // For now, if there's no data, the schema might be incomplete or incorrect for ColNameExpr. + // LiteralExprs can still define their part of the schema. + if finalSchema == nil { // if t.data was empty + newSchemaSeries = make([]df.SeriesSchema, 0, len(expressions)) + for _, expr := range expressions { + outputName := expr.Name() + switch e := expr.(type) { + case df.ColNameExpr: + // Cannot determine format without data or schema introspection of original df + // This part needs refinement: how to get schema for a col if no data? + // Fallback: try to get from original schema if possible + originalSeriesSchema, found := t.schema.GetByName(e.Col()) + if !found { + panic(fmt.Sprintf("cannot determine schema for column %s with no data and not in original schema", e.Col())) + } + if outputName == "" { outputName = e.Col() } + newSchemaSeries = append(newSchemaSeries, df.SeriesSchema{Name: outputName, Format: originalSeriesSchema.Format}) + case df.LiteralExpr: + litVal := e.Const() + if outputName == "" { panic(fmt.Sprintf("literal expression %v must have an alias", e.Const().GetAsString()))} + newSchemaSeries = append(newSchemaSeries, df.SeriesSchema{Name: outputName, Format: litVal.Schema()}) + default: + panic(fmt.Sprintf("unsupported expression type for schema generation with no data: %T", expr)) + } + } + finalSchema = NewInMemorySchema(t.name+"_select", newSchemaSeries) + } + } + + + return NewDataframeFromRowAndName(t.name+"_select", finalSchema, &newRowsData) } func (t *inmemoryDataFrame) SelectBySeriesIndex(index ...int) (d df.DataFrame) { @@ -439,7 +536,7 @@ func (t *inmemoryDataFrame) Append(d df.DataFrame) df.DataFrame { return NewDataframeFromRow(t.schema, &s1) } -func (t *inmemoryDataFrame) Group(others ...string) df.GroupedDataFrame { +func (t *inmemoryDataFrame) GroupBy(others ...string) df.GroupedDataFrame { return NewGroupedDf(t, others...) } diff --git a/df/inmemory/df_benchmark_test.go b/df/inmemory/df_benchmark_test.go index c3a3b93..9467357 100644 --- a/df/inmemory/df_benchmark_test.go +++ b/df/inmemory/df_benchmark_test.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/df_test.go b/df/inmemory/df_test.go index 88fa328..9d8237b 100644 --- a/df/inmemory/df_test.go +++ b/df/inmemory/df_test.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( @@ -99,8 +100,8 @@ func TestInMemoryDf(t *testing.T) { assert.Equal(t, "renamed", data.Name()) assert.Equal(t, "renamed", renamedData.Name()) - // Group - grouped := data.Group("c1") + // GroupBy + grouped := data.GroupBy("c1") assert.Equal(t, int64(4), grouped.Len()) // Append diff --git a/df/inmemory/expr.go b/df/inmemory/expr.go index 524eb4e..cf80c5d 100644 --- a/df/inmemory/expr.go +++ b/df/inmemory/expr.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/grouped_df.go b/df/inmemory/grouped_df.go index 9f1ea5a..c9489ef 100644 --- a/df/inmemory/grouped_df.go +++ b/df/inmemory/grouped_df.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( @@ -62,6 +63,290 @@ func (t *inmemoryGroupedDataFrame) Len() int64 { return int64(len(t.data)) } +// Agg performs aggregations on the groups. +func (t *inmemoryGroupedDataFrame) Agg(configs ...df.AggregationConfig) df.DataFrame { + if len(t.data) == 0 && len(configs) == 0 { + return NewDataframe(NewInMemorySchema("", []df.SeriesSchema{})) + } + + // Determine output schema + var outputSeriesSchemas []df.SeriesSchema + var keySchema df.DataFrameSchema + if len(t.keys) > 0 { + // Get schema from the first key (all keys have the same schema) + for _, kRow := range t.keys { + keySchema = kRow.Schema() + break + } + for i := 0; i < keySchema.Len(); i++ { + outputSeriesSchemas = append(outputSeriesSchemas, keySchema.Get(i)) + } + } + + for _, aggConfig := range configs { + // TODO: Determine actual type based on Func and InputCol type + // For now, count is Int64, others could be Float64 or original type + var aggSchema df.SeriesSchema + switch strings.ToLower(aggConfig.Func) { + case "count": + aggSchema = df.SeriesSchema{Name: aggConfig.OutputColName, Format: int64Format} + case "sum", "mean", "min", "max": + // Placeholder: This needs to be more robust, inspect input column type + // For simplicity, assume float64 for mean, or try to match input type for others. + // This part will be very complex in a full implementation. + if keySchema != nil && aggConfig.InputCol != "" { + inputColSchema, ok := keySchema.GetByName(aggConfig.InputCol) + if !ok && len(t.data) > 0 { + // Fallback: check the schema of the first group's data + for _, groupDf := range t.data { + inputColSchema, ok = groupDf.Schema().GetByName(aggConfig.InputCol) + if ok { + break + } + } + } + + if ok { + if strings.ToLower(aggConfig.Func) == "mean" { + aggSchema = df.SeriesSchema{Name: aggConfig.OutputColName, Format: float64Format} + } else { + aggSchema = df.SeriesSchema{Name: aggConfig.OutputColName, Format: inputColSchema.Format} + } + } else { + // Could not determine input type, default to float64 or error + // For now, let's use float64 as a default for non-count if input col is specified + aggSchema = df.SeriesSchema{Name: aggConfig.OutputColName, Format: float64Format} + } + } else { + // Default for sum, min, max if InputCol is missing (though usually required) + aggSchema = df.SeriesSchema{Name: aggConfig.OutputColName, Format: float64Format} + } + default: + panic(fmt.Sprintf("unsupported aggregation function: %s", aggConfig.Func)) + } + outputSeriesSchemas = append(outputSeriesSchemas, aggSchema) + } + outputSchema := NewInMemorySchema("", outputSeriesSchemas) + outputRows := make([]df.Row, 0, len(t.data)) + + for keyStr, groupDf := range t.data { + keyRow := t.keys[keyStr] + newRowValues := make([]df.Value, 0, outputSchema.Len()) + + // Add key values + for i := 0; i < keyRow.Len(); i++ { + newRowValues = append(newRowValues, keyRow.Get(i)) + } + + // Calculate aggregations + for _, aggConfig := range configs { + var aggValue df.Value + inputColName := aggConfig.InputCol + aggFunc := strings.ToLower(aggConfig.Func) + + switch aggFunc { + case "count": + if inputColName == "" { + aggValue = NewInt64Value(groupDf.Len()) + } else { + colIdx := groupDf.Schema().GetIndexByName(inputColName) + if colIdx == -1 { + panic(fmt.Sprintf("column %s not found for count", inputColName)) + } + count := int64(0) + groupDf.ForEachRow(func(r df.Row) { + if !r.IsNil(colIdx) { + count++ + } + }) + aggValue = NewInt64Value(count) + } + case "sum": + colIdx := groupDf.Schema().GetIndexByName(inputColName) + if colIdx == -1 { + panic(fmt.Sprintf("column %s not found for sum", inputColName)) + } + // This is a simplified sum, assuming numeric types that can be summed. + // A full implementation needs type checking and dispatch. + currentSum := 0.0 // Default to float64 for sum for now + isFirst := true + var sumVal df.Value + + groupDf.ForEachRow(func(r df.Row) { + val := r.Get(colIdx) + if val.IsNil() { + return + } + // Basic sum for float64, int64. Others would need conversion or type assertion. + switch v := val.Get().(type) { + case float64: + if isFirst { currentSum = 0; isFirst = false; sumVal = NewFloat64Value(0)} + currentSum += v + sumVal = NewFloat64Value(currentSum) + case int64: + if isFirst { currentSum = 0; isFirst = false; sumVal = NewInt64Value(0)} + currentSum += float64(v) // Promote to float for generic sum + // Determine if original type was int to store back as int if possible + // For now, sum promotes to float64 + sumVal = NewFloat64Value(currentSum) + + case int: + if isFirst { currentSum = 0; isFirst = false; sumVal = NewInt64Value(0)} + currentSum += float64(v) + sumVal = NewFloat64Value(currentSum) + default: + // Try to convert to float64 if possible + fv, err := val.Schema().Convert(val.Get()) + if err == nil { + if fvFloat, ok := fv.(float64); ok { + if isFirst { currentSum = 0; isFirst = false; sumVal = NewFloat64Value(0) } + currentSum += fvFloat + sumVal = NewFloat64Value(currentSum) + return + } + } + panic(fmt.Sprintf("unsupported type for sum: %T on column %s", val.Get(), inputColName)) + } + }) + if isFirst { // No non-nil values + // Find the type of the column for nil value + seriesSchema := outputSchema.GetByName(aggConfig.OutputColName) + if seriesSchema.Format.Type() == float64Format.Type() { + aggValue = NewFloat64Nil() + } else if seriesSchema.Format.Type() == int64Format.Type() { + aggValue = NewInt64Nil() + } else { + // Default or panic for unsupported type for nil + aggValue = NewFloat64Nil() // Fallback + } + } else { + aggValue = sumVal + } + + case "mean": + colIdx := groupDf.Schema().GetIndexByName(inputColName) + if colIdx == -1 { + panic(fmt.Sprintf("column %s not found for mean", inputColName)) + } + sum := 0.0 + count := int64(0) + groupDf.ForEachRow(func(r df.Row) { + val := r.Get(colIdx) + if !val.IsNil() { + // Assuming numeric types convertible to float64 + // A robust solution would check types and handle errors + floatVal, err := val.Schema().Convert(val.Get()) + if err == nil { + switch v := floatVal.(type) { + case float64: sum += v + case int64: sum += float64(v) + case int: sum += float64(v) + default: + panic(fmt.Sprintf("unsupported type for mean: %T on column %s after conversion", v, inputColName)) + } + count++ + } else { + panic(fmt.Sprintf("cannot convert value for mean on column %s: %v", inputColName, err)) + } + } + }) + if count == 0 { + aggValue = NewFloat64Nil() // Or handle as error / NaN + } else { + aggValue = NewFloat64Value(sum / float64(count)) + } + case "min", "max": + colIdx := groupDf.Schema().GetIndexByName(inputColName) + if colIdx == -1 { + panic(fmt.Sprintf("column %s not found for %s", inputColName, aggFunc)) + } + var resVal df.Value = nil + groupDf.ForEachRow(func(r df.Row) { + val := r.Get(colIdx) + if val.IsNil() { + return + } + if resVal == nil || resVal.IsNil() { + resVal = val + return + } + + // This requires comparable types. + // Simplified comparison logic. Real implementation needs type-specific comparisons. + currentFloat, cfOk := convertToFloatForCompare(resVal) + newFloat, nfOk := convertToFloatForCompare(val) + + if cfOk && nfOk { + if aggFunc == "min" { + if newFloat < currentFloat { + resVal = val + } + } else { // max + if newFloat > currentFloat { + resVal = val + } + } + } else { + // Fallback to string comparison or error if types are not directly comparable as numbers + // This is a placeholder for more robust type handling + sCurrent := resVal.GetAsString() + sNew := val.GetAsString() + if aggFunc == "min" { + if strings.Compare(sNew, sCurrent) < 0 { + resVal = val + } + } else { // max + if strings.Compare(sNew, sCurrent) > 0 { + resVal = val + } + } + } + }) + + if resVal == nil { // Group was empty or all values were nil + // Determine the correct nil type based on the output schema for this agg column + seriesSchema := outputSchema.GetByName(aggConfig.OutputColName) + // This is a simplification. Ideally, df.Value itself should have a typed Nil constructor or similar. + if seriesSchema.Format.Name() == "float64" { aggValue = NewFloat64Nil() } else + if seriesSchema.Format.Name() == "int64" { aggValue = NewInt64Nil() } else + if seriesSchema.Format.Name() == "string" { aggValue = NewStringNil() } else + { aggValue = NewFloat64Nil() /* Default nil type */ } + + } else { + aggValue = resVal + } + + default: + panic(fmt.Sprintf("unsupported aggregation function: %s", aggConfig.Func)) + } + newRowValues = append(newRowValues, aggValue) + } + outputRows = append(outputRows, NewInMemoryRow(outputSchema, newRowValues)) + } + + return NewDataframeFromRow(outputSchema, &outputRows) +} + +// Helper function for min/max comparison, tries to convert to float64 +// This is a simplification. A full solution needs proper type handling and comparison. +func convertToFloatForCompare(v df.Value) (float64, bool) { + if v.IsNil() { + return 0, false + } + switch val := v.Get().(type) { + case float64: + return val, true + case int64: + return float64(val), true + case int: + return float64(val), true + // Add other numeric types if necessary + default: + return 0, false // Cannot convert to float64 for comparison + } +} + + func getKey(r df.Row) string { var b strings.Builder for i := 0; i < r.Len(); i++ { diff --git a/df/inmemory/grouped_df_test.go b/df/inmemory/grouped_df_test.go index ff3f4df..4f5ab30 100644 --- a/df/inmemory/grouped_df_test.go +++ b/df/inmemory/grouped_df_test.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/grouped_series.go b/df/inmemory/grouped_series.go index 022c847..b6a6005 100644 --- a/df/inmemory/grouped_series.go +++ b/df/inmemory/grouped_series.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/grouped_series_test.go b/df/inmemory/grouped_series_test.go index c6f4a8d..df17460 100644 --- a/df/inmemory/grouped_series_test.go +++ b/df/inmemory/grouped_series_test.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/merged_df.go b/df/inmemory/merged_df.go index a1f7480..51ee6c9 100644 --- a/df/inmemory/merged_df.go +++ b/df/inmemory/merged_df.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/merged_df_test.go b/df/inmemory/merged_df_test.go index d4abef0..5074785 100644 --- a/df/inmemory/merged_df_test.go +++ b/df/inmemory/merged_df_test.go @@ -1 +1,251 @@ -package inmemory +//go:build inmemory + +package inmemory_test // Test files should be in _test package + +import ( + "fmt" + "testing" + + "github.com/blue4209211/pq/df" + "github.com/blue4209211/pq/df/inmemory" + "github.com/stretchr/testify/assert" +) + +// Helper to create an inmemory.DataFrame for testing +func createTestInMemoryDF(t *testing.T, name string, colDefs []df.SeriesSchema, rowsData [][]interface{}) df.DataFrame { + t.Helper() + schema := df.NewSchema(name, colDefs) + + rows := make([]df.Row, len(rowsData)) + for i, rowData := range rowsData { + values := make([]df.Value, len(rowData)) + if len(rowData) != len(colDefs) { + t.Fatalf("Row data length %d does not match colDefs length %d for row %d", len(rowData), len(colDefs), i) + } + for j, cellData := range rowData { + colFormat := colDefs[j].Format + // Use concrete types from inmemory package + switch colFormat.Name() { + case df.IntegerFormat.Name(): + if cellData == nil { + values[j] = inmemory.NewIntValue(nil) + } else { + values[j] = inmemory.NewIntValueConst(cellData.(int64)) + } + case df.StringFormat.Name(): + if cellData == nil { + values[j] = inmemory.NewStringValue(nil) + } else { + values[j] = inmemory.NewStringValueConst(cellData.(string)) + } + case df.DoubleFormat.Name(): + if cellData == nil { + values[j] = inmemory.NewDoubleValue(nil) + } else { + values[j] = inmemory.NewDoubleValueConst(cellData.(float64)) + } + case df.BoolFormat.Name(): + if cellData == nil { + values[j] = inmemory.NewBoolValue(nil) + } else { + values[j] = inmemory.NewBoolValueConst(cellData.(bool)) + } + default: + t.Fatalf("Unsupported format in test helper: %s", colFormat.Name()) + } + } + rows[i] = inmemory.NewRow(&schema, &values) + } + return inmemory.NewDataframeFromRowAndName(name, schema, &rows) +} + +// Helper to convert DataFrame to slice of slices for easier comparison +func dfToSlice(t *testing.T, dataFrame df.DataFrame) [][]interface{} { + t.Helper() + var result [][]interface{} + if dataFrame == nil { + return result + } + for r := int64(0); r < dataFrame.Len(); r++ { + row := dataFrame.GetRow(r) + var rowData []interface{} + for c := 0; c < row.Len(); c++ { + val := row.Get(c) + if val.IsNil() { + rowData = append(rowData, nil) // Use actual nil for easier comparison with expected + } else { + rowData = append(rowData, val.Get()) + } + } + result = append(result, rowData) + } + return result +} + + +func TestNewMergeDataframe_EmptyInput(t *testing.T) { + mergedDf, err := inmemory.NewMergeDataframe("test_empty") + assert.Error(t, err, "Expected error for empty input") + if err != nil { // Check error message only if error is not nil + assert.Equal(t, "empty data", err.Error()) + } + assert.Nil(t, mergedDf, "DataFrame should be nil on error") +} + +func TestNewMergeDataframe_SingleDataFrame(t *testing.T) { + cols := []df.SeriesSchema{ + {Name: "id", Format: df.IntegerFormat}, + {Name: "name", Format: df.StringFormat}, + } + data1 := [][]interface{}{ + {int64(1), "Alice"}, + {int64(2), "Bob"}, + } + df1 := createTestInMemoryDF(t, "df1", cols, data1) + + mergedDf, err := inmemory.NewMergeDataframe("renamed_df", df1) + assert.NoError(t, err) + assert.NotNil(t, mergedDf) + + assert.Equal(t, "renamed_df", mergedDf.Name(), "Merged DF name should be updated") + assert.Equal(t, "df1", df1.Name(), "Original DF name should be unchanged") + + assert.True(t, df1.Schema().Equals(mergedDf.Schema()), "Schemas should be equal") + assert.Equal(t, df1.Len(), mergedDf.Len(), "Lengths should be equal") + assert.Equal(t, dfToSlice(t, df1), dfToSlice(t, mergedDf), "Row data should be identical") +} + +func TestNewMergeDataframe_Multiple_IdenticalSchemas(t *testing.T) { + cols := []df.SeriesSchema{ + {Name: "id", Format: df.IntegerFormat}, + {Name: "value", Format: df.StringFormat}, + } + data1 := [][]interface{}{{int64(1), "A"}, {int64(2), "B"}} + df1 := createTestInMemoryDF(t, "df1", cols, data1) + + data2 := [][]interface{}{{int64(3), "C"}, {int64(4), "D"}} + df2 := createTestInMemoryDF(t, "df2", cols, data2) + + data3 := [][]interface{}{{int64(5), "E"}} + df3 := createTestInMemoryDF(t, "df3", cols, data3) + + mergedDf, err := inmemory.NewMergeDataframe("merged_identical", df1, df2, df3) + assert.NoError(t, err) + assert.NotNil(t, mergedDf) + + assert.Equal(t, "merged_identical", mergedDf.Name()) + assert.True(t, df1.Schema().Equals(mergedDf.Schema()), "Schema should be from the first DataFrame") + expectedLen := df1.Len() + df2.Len() + df3.Len() + assert.Equal(t, expectedLen, mergedDf.Len(), "Length should be sum of input lengths") + + expectedData := [][]interface{}{ + {int64(1), "A"}, {int64(2), "B"}, + {int64(3), "C"}, {int64(4), "D"}, + {int64(5), "E"}, + } + assert.Equal(t, expectedData, dfToSlice(t, mergedDf), "Row data should be concatenated") +} + +func TestNewMergeDataframe_WithEmptyDataFrames(t *testing.T) { + cols := []df.SeriesSchema{ + {Name: "val", Format: df.DoubleFormat}, + } + data1 := [][]interface{}{{1.1}, {2.2}} + df1 := createTestInMemoryDF(t, "df1", cols, data1) + + dfEmpty := createTestInMemoryDF(t, "dfEmpty", cols, [][]interface{}{}) + + data2 := [][]interface{}{{3.3}} + df2 := createTestInMemoryDF(t, "df2", cols, data2) + + mergedDf1, err1 := inmemory.NewMergeDataframe("merged_with_empty1", df1, dfEmpty, df2) + assert.NoError(t, err1) + assert.NotNil(t, mergedDf1) + assert.Equal(t, df1.Len()+df2.Len(), mergedDf1.Len()) + assert.True(t, df1.Schema().Equals(mergedDf1.Schema())) + expectedData1 := [][]interface{}{{1.1}, {2.2}, {3.3}} + assert.Equal(t, expectedData1, dfToSlice(t, mergedDf1)) + + mergedDf2, err2 := inmemory.NewMergeDataframe("merged_with_empty2", dfEmpty, df1, df2) + assert.NoError(t, err2) + assert.NotNil(t, mergedDf2) + assert.Equal(t, df1.Len()+df2.Len(), mergedDf2.Len()) + assert.True(t, dfEmpty.Schema().Equals(mergedDf2.Schema())) + expectedData2 := [][]interface{}{{1.1}, {2.2}, {3.3}} + assert.Equal(t, expectedData2, dfToSlice(t, mergedDf2)) + + dfEmpty2 := createTestInMemoryDF(t, "dfEmpty2", cols, [][]interface{}{}) + mergedDf3, err3 := inmemory.NewMergeDataframe("merged_all_empty", dfEmpty, dfEmpty2) + assert.NoError(t, err3) + assert.NotNil(t, mergedDf3) + assert.Equal(t, int64(0), mergedDf3.Len()) + assert.True(t, dfEmpty.Schema().Equals(mergedDf3.Schema())) +} + + +func TestNewMergeDataframe_SchemaCompatibility_Implicit(t *testing.T) { + t.Run("CompatibleStructureDifferentNames", func(t *testing.T) { + cols1 := []df.SeriesSchema{ + {Name: "colA", Format: df.IntegerFormat}, + {Name: "colB", Format: df.StringFormat}, + } + data1 := [][]interface{}{{int64(1), "hello"}} + df1 := createTestInMemoryDF(t, "df1_names", cols1, data1) + + cols2 := []df.SeriesSchema{ + {Name: "fieldA", Format: df.IntegerFormat}, + {Name: "fieldB", Format: df.StringFormat}, + } + data2 := [][]interface{}{{int64(100), "world"}} + df2 := createTestInMemoryDF(t, "df2_names", cols2, data2) + + mergedDf, err := inmemory.NewMergeDataframe("merged_compat_names", df1, df2) + assert.NoError(t, err) + assert.NotNil(t, mergedDf) + + assert.Equal(t, "colA", mergedDf.Schema().Get(0).Name) + assert.Equal(t, "colB", mergedDf.Schema().Get(1).Name) + + expectedData := [][]interface{}{ + {int64(1), "hello"}, + {int64(100), "world"}, + } + assert.Equal(t, expectedData, dfToSlice(t, mergedDf)) + }) + + t.Run("IncompatibleTypes", func(t *testing.T) { + cols1 := []df.SeriesSchema{ + {Name: "colA", Format: df.IntegerFormat}, + {Name: "colB", Format: df.StringFormat}, + } + data1 := [][]interface{}{{int64(1), "hello"}} + df1 := createTestInMemoryDF(t, "df1_types", cols1, data1) + + cols2 := []df.SeriesSchema{ + {Name: "colA", Format: df.StringFormat}, + {Name: "colB", Format: df.IntegerFormat}, + } + data2 := [][]interface{}{{"test", int64(99)}} + df2 := createTestInMemoryDF(t, "df2_types", cols2, data2) + + mergedDf, err := inmemory.NewMergeDataframe("merged_incompat_types", df1, df2) + assert.NoError(t, err, "NewMergeDataframe itself does not error on schema type incompatibility") + assert.NotNil(t, mergedDf) + + row0 := mergedDf.GetRow(0) + assert.Equal(t, int64(1), row0.Get(0).GetAsInt()) + assert.Equal(t, "hello", row0.Get(1).GetAsString()) + + row1 := mergedDf.GetRow(1) + // df1 schema: colA (idx 0) is int, colB (idx 1) is string + // df2 data was: "test" (for colA's position), int64(99) (for colB's position) + + // The inmemory.Value concrete types will be StringValue and IntValue. + // Accessing them with the wrong GetAs() will panic. + assert.Panics(t, func() { _ = row1.Get(0).GetAsInt() }, "Accessing string as int should panic") + assert.Equal(t, "test", row1.Get(0).GetAsString(), "Accessing string as string should work") + + assert.Panics(t, func() { _ = row1.Get(1).GetAsString() }, "Accessing int as string should panic") + assert.Equal(t, int64(99), row1.Get(1).GetAsInt(), "Accessing int as int should work") + }) +} diff --git a/df/inmemory/row.go b/df/inmemory/row.go index 4bdae7a..20642f1 100644 --- a/df/inmemory/row.go +++ b/df/inmemory/row.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/row_test.go b/df/inmemory/row_test.go index f3a3477..efbf870 100644 --- a/df/inmemory/row_test.go +++ b/df/inmemory/row_test.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/series.go b/df/inmemory/series.go index f3c4d87..ee3fc3d 100644 --- a/df/inmemory/series.go +++ b/df/inmemory/series.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/series_benchmark_test.go b/df/inmemory/series_benchmark_test.go index bf3f72e..aa55707 100644 --- a/df/inmemory/series_benchmark_test.go +++ b/df/inmemory/series_benchmark_test.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/series_bool.go b/df/inmemory/series_bool.go index 5338b88..5849bfa 100644 --- a/df/inmemory/series_bool.go +++ b/df/inmemory/series_bool.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/series_bool_test.go b/df/inmemory/series_bool_test.go index dc5858c..ed5f998 100644 --- a/df/inmemory/series_bool_test.go +++ b/df/inmemory/series_bool_test.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/series_datetime.go b/df/inmemory/series_datetime.go index dfb14fc..453a22c 100644 --- a/df/inmemory/series_datetime.go +++ b/df/inmemory/series_datetime.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/series_datetime_test.go b/df/inmemory/series_datetime_test.go index 5495500..098e45d 100644 --- a/df/inmemory/series_datetime_test.go +++ b/df/inmemory/series_datetime_test.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/series_float64.go b/df/inmemory/series_float64.go index e6abe29..d1ec1a5 100644 --- a/df/inmemory/series_float64.go +++ b/df/inmemory/series_float64.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/series_float64_test.go b/df/inmemory/series_float64_test.go index 215961d..06c5a71 100644 --- a/df/inmemory/series_float64_test.go +++ b/df/inmemory/series_float64_test.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/series_int64.go b/df/inmemory/series_int64.go index a624a7c..0e597b2 100644 --- a/df/inmemory/series_int64.go +++ b/df/inmemory/series_int64.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/series_int64_test.go b/df/inmemory/series_int64_test.go index b0e0be8..9ba9be1 100644 --- a/df/inmemory/series_int64_test.go +++ b/df/inmemory/series_int64_test.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/series_string.go b/df/inmemory/series_string.go index 8d9fc42..01f206e 100644 --- a/df/inmemory/series_string.go +++ b/df/inmemory/series_string.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/series_string_test.go b/df/inmemory/series_string_test.go index 824cc20..66148b2 100644 --- a/df/inmemory/series_string_test.go +++ b/df/inmemory/series_string_test.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/series_test.go b/df/inmemory/series_test.go index 102f240..2c6f0ac 100644 --- a/df/inmemory/series_test.go +++ b/df/inmemory/series_test.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/series_val.go b/df/inmemory/series_val.go index 54cf72f..ed9314d 100644 --- a/df/inmemory/series_val.go +++ b/df/inmemory/series_val.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( diff --git a/df/inmemory/series_val_test.go b/df/inmemory/series_val_test.go index 00297ed..d97494c 100644 --- a/df/inmemory/series_val_test.go +++ b/df/inmemory/series_val_test.go @@ -1,3 +1,4 @@ +//go:build inmemory package inmemory import ( @@ -42,8 +43,9 @@ func TestNewDataFrameVal(t *testing.T) { assert.Equal(t, df.StringFormat, val.Schema()) assert.Equal(t, true, val.IsNil()) assert.Equal(t, nil, val.Get()) - //TODO assert panic - //assert.Equal(t, "", val.GetAsString()) + assert.PanicsWithValue(t, "GetAsString() called on a nil value", func() { + val.GetAsString() + }, "Calling GetAsString() on a nil string value should panic with the specified message.") } func TestNewDataFrameValEqual(t *testing.T) { diff --git a/go.mod b/go.mod index 6478785..c6471da 100644 --- a/go.mod +++ b/go.mod @@ -3,7 +3,7 @@ module github.com/blue4209211/pq go 1.18 require ( - cloud.google.com/go/storage v1.27.0 + cloud.google.com/go/storage v1.30.1 github.com/apache/arrow/go/v7 v7.0.1 github.com/aws/aws-sdk-go v1.44.131 github.com/go-sql-driver/mysql v1.6.0 @@ -13,53 +13,54 @@ require ( github.com/mattn/go-sqlite3 v1.14.16 github.com/olekukonko/tablewriter v0.0.5 github.com/samber/lo v1.33.0 - github.com/stretchr/testify v1.8.1 + github.com/stretchr/testify v1.8.4 github.com/xo/dburl v0.12.4 - golang.org/x/exp v0.0.0-20221031165847-c99f073a8326 - google.golang.org/api v0.102.0 + golang.org/x/exp v0.0.0-20231006140011-7918f672742d + google.golang.org/api v0.126.0 ) require ( - cloud.google.com/go v0.105.0 // indirect - cloud.google.com/go/compute v1.12.1 // indirect - cloud.google.com/go/compute/metadata v0.2.1 // indirect - cloud.google.com/go/iam v0.7.0 // indirect + cloud.google.com/go v0.110.4 // indirect + cloud.google.com/go/compute v1.21.0 // indirect + cloud.google.com/go/compute/metadata v0.2.3 // indirect + cloud.google.com/go/iam v1.1.1 // indirect github.com/JohnCGriffin/overflow v0.0.0-20211019200055-46fa312c352c // indirect - github.com/andybalholm/brotli v1.0.4 // indirect + github.com/andybalholm/brotli v1.0.5 // indirect + github.com/apache/arrow/go/v14 v14.0.2 // indirect github.com/apache/thrift v0.17.0 // indirect github.com/davecgh/go-spew v1.1.1 // indirect github.com/dimchansky/utfbom v1.1.1 // indirect - github.com/goccy/go-json v0.9.11 // indirect + github.com/goccy/go-json v0.10.2 // indirect github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect - github.com/golang/protobuf v1.5.2 // indirect - github.com/google/flatbuffers v22.10.26+incompatible // indirect + github.com/golang/protobuf v1.5.3 // indirect + github.com/google/flatbuffers v23.5.26+incompatible // indirect github.com/google/go-cmp v0.5.9 // indirect - github.com/google/uuid v1.3.0 // indirect - github.com/googleapis/enterprise-certificate-proxy v0.2.0 // indirect - github.com/googleapis/gax-go/v2 v2.7.0 // indirect + github.com/google/uuid v1.3.1 // indirect + github.com/googleapis/enterprise-certificate-proxy v0.2.3 // indirect + github.com/googleapis/gax-go/v2 v2.11.0 // indirect github.com/jmespath/go-jmespath v0.4.0 // indirect github.com/klauspost/asmfmt v1.3.2 // indirect - github.com/klauspost/compress v1.15.12 // indirect - github.com/klauspost/cpuid/v2 v2.1.2 // indirect + github.com/klauspost/compress v1.16.7 // indirect + github.com/klauspost/cpuid/v2 v2.2.5 // indirect github.com/mattn/go-runewidth v0.0.14 // indirect github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 // indirect github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 // indirect - github.com/pierrec/lz4/v4 v4.1.17 // indirect + github.com/pierrec/lz4/v4 v4.1.18 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/rivo/uniseg v0.4.2 // indirect github.com/zeebo/xxh3 v1.0.2 // indirect go.opencensus.io v0.24.0 // indirect - golang.org/x/mod v0.6.0 // indirect - golang.org/x/net v0.1.0 // indirect - golang.org/x/oauth2 v0.1.0 // indirect - golang.org/x/sys v0.1.0 // indirect - golang.org/x/text v0.4.0 // indirect - golang.org/x/tools v0.2.0 // indirect + golang.org/x/mod v0.13.0 // indirect + golang.org/x/net v0.17.0 // indirect + golang.org/x/oauth2 v0.10.0 // indirect + golang.org/x/sys v0.13.0 // indirect + golang.org/x/text v0.13.0 // indirect + golang.org/x/tools v0.14.0 // indirect golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 // indirect google.golang.org/appengine v1.6.7 // indirect - google.golang.org/genproto v0.0.0-20221027153422-115e99e71e1c // indirect - google.golang.org/grpc v1.50.1 // indirect - google.golang.org/protobuf v1.28.1 // indirect + google.golang.org/genproto v0.0.0-20230711160842-782d3b101e98 // indirect + google.golang.org/grpc v1.58.2 // indirect + google.golang.org/protobuf v1.31.0 // indirect gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect gopkg.in/yaml.v2 v2.4.0 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect diff --git a/go.sum b/go.sum index c7ac286..ad5ead5 100644 --- a/go.sum +++ b/go.sum @@ -2,15 +2,26 @@ cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMT cloud.google.com/go v0.34.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw= cloud.google.com/go v0.105.0 h1:DNtEKRBAAzeS4KyIory52wWHuClNaXJ5x1F7xa4q+5Y= cloud.google.com/go v0.105.0/go.mod h1:PrLgOJNe5nfE9UMxKxgXj4mD3voiP+YQ6gdt6KMFOKM= +cloud.google.com/go v0.110.4 h1:1JYyxKMN9hd5dR2MYTPWkGUgcoxVVhg0LKNKEo0qvmk= +cloud.google.com/go v0.110.4/go.mod h1:+EYjdK8e5RME/VY/qLCAtuyALQ9q67dvuum8i+H5xsI= cloud.google.com/go/compute v1.12.1 h1:gKVJMEyqV5c/UnpzjjQbo3Rjvvqpr9B1DFSbJC4OXr0= cloud.google.com/go/compute v1.12.1/go.mod h1:e8yNOBcBONZU1vJKCvCoDw/4JQsA0dpM4x/6PIIOocU= +cloud.google.com/go/compute v1.21.0 h1:JNBsyXVoOoNJtTQcnEY5uYpZIbeCTYIeDe0Xh1bySMk= +cloud.google.com/go/compute v1.21.0/go.mod h1:4tCnrn48xsqlwSAiLf1HXMQk8CONslYbdiEZc9FEIbM= cloud.google.com/go/compute/metadata v0.2.1 h1:efOwf5ymceDhK6PKMnnrTHP4pppY5L22mle96M1yP48= cloud.google.com/go/compute/metadata v0.2.1/go.mod h1:jgHgmJd2RKBGzXqF5LR2EZMGxBkeanZ9wwa75XHJgOM= +cloud.google.com/go/compute/metadata v0.2.3 h1:mg4jlk7mCAj6xXp9UJ4fjI9VUI5rubuGBW5aJ7UnBMY= +cloud.google.com/go/compute/metadata v0.2.3/go.mod h1:VAV5nSsACxMJvgaAuX6Pk2AawlZn8kiOGuCv6gTkwuA= cloud.google.com/go/iam v0.7.0 h1:k4MuwOsS7zGJJ+QfZ5vBK8SgHBAvYN/23BWsiihJ1vs= cloud.google.com/go/iam v0.7.0/go.mod h1:H5Br8wRaDGNc8XP3keLc4unfUUZeyH3Sfl9XpQEYOeg= +cloud.google.com/go/iam v1.1.1 h1:lW7fzj15aVIXYHREOqjRBV9PsH0Z6u8Y46a1YGvQP4Y= +cloud.google.com/go/iam v1.1.1/go.mod h1:A5avdyVL2tCppe4unb0951eI9jreack+RJ0/d+KUZOU= cloud.google.com/go/longrunning v0.1.1 h1:y50CXG4j0+qvEukslYFBCrzaXX0qpFbBzc3PchSu/LE= +cloud.google.com/go/longrunning v0.5.1 h1:Fr7TXftcqTudoyRJa113hyaqlGdiBQkp0Gq7tErFDWI= cloud.google.com/go/storage v1.27.0 h1:YOO045NZI9RKfCj1c5A/ZtuuENUc8OAW+gHdGnDgyMQ= cloud.google.com/go/storage v1.27.0/go.mod h1:x9DOL8TK/ygDUMieqwfhdpQryTeEkhGKMi80i/iqR2s= +cloud.google.com/go/storage v1.30.1 h1:uOdMxAs8HExqBlnLtnQyP0YkvbiDpdGShGKtx6U/oNM= +cloud.google.com/go/storage v1.30.1/go.mod h1:NfxhC0UJE1aXSx7CIIbCf7y9HKT7BiccwkR7+P7gN8E= dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU= dmitri.shuralyov.com/gpu/mtl v0.0.0-20201218220906-28db891af037/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU= gioui.org v0.0.0-20210308172011-57750fc8a0a6/go.mod h1:RSH6KIUZ0p2xy5zHDxgAM4zumjgTw83q2ge/PI+yyw8= @@ -31,7 +42,11 @@ github.com/alecthomas/units v0.0.0-20190717042225-c3de453c63f4/go.mod h1:ybxpYRF github.com/andybalholm/brotli v1.0.3/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig= github.com/andybalholm/brotli v1.0.4 h1:V7DdXeJtZscaqfNuAdSRuRFzuiKlHSC/Zh3zl9qY3JY= github.com/andybalholm/brotli v1.0.4/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig= +github.com/andybalholm/brotli v1.0.5 h1:8uQZIdzKmjc/iuPu7O2ioW48L81FgatrcpfFmiq/cCs= +github.com/andybalholm/brotli v1.0.5/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig= github.com/antihax/optional v1.0.0/go.mod h1:uupD/76wgC+ih3iEmQUL+0Ugr19nfwCT1kdvxnR2qWY= +github.com/apache/arrow/go/v14 v14.0.2 h1:N8OkaJEOfI3mEZt07BIkvo4sC6XDbL+48MBPWO5IONw= +github.com/apache/arrow/go/v14 v14.0.2/go.mod h1:u3fgh3EdgN/YQ8cVQRguVW3R+seMybFg8QBQ5LU+eBY= github.com/apache/arrow/go/v7 v7.0.1 h1:WpCfq+AQxvXaI6/KplHE27MPMFx5av0o5NbPCTAGfy4= github.com/apache/arrow/go/v7 v7.0.1/go.mod h1:JxDpochJbCVxqbX4G8i1jRqMrnTCQdf8pTccAfLD8Es= github.com/apache/thrift v0.12.0/go.mod h1:cp2SuWMxlEZw2r+iP2GNCdIi4C1qmUzdZFSVb+bacwQ= @@ -119,6 +134,8 @@ github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/me github.com/goccy/go-json v0.7.10/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I= github.com/goccy/go-json v0.9.11 h1:/pAaQDLHEoCq/5FFmSKBswWmK6H0e8g4159Kc/X/nqk= github.com/goccy/go-json v0.9.11/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I= +github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU= +github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I= github.com/gogo/googleapis v1.1.0/go.mod h1:gf4bu3Q80BeJ6H1S1vYPm8/ELATdvryBaNFGgqEef3s= github.com/gogo/protobuf v1.1.1/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ= github.com/gogo/protobuf v1.2.0/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ= @@ -147,6 +164,8 @@ github.com/golang/protobuf v1.4.3/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= github.com/golang/protobuf v1.5.2 h1:ROPKBNFfQgOUMifHyP+KYbvpjbdoFNs+aK7DXlji0Tw= github.com/golang/protobuf v1.5.2/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY= +github.com/golang/protobuf v1.5.3 h1:KhyjKVUg7Usr/dYsdSqoFveMYd5ko72D+zANwlG1mmg= +github.com/golang/protobuf v1.5.3/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY= github.com/golang/snappy v0.0.0-20180518054509-2e65f85255db/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q= github.com/golang/snappy v0.0.4 h1:yAGX7huGHXlcLOEtBnF4w7FQwA26wojNCwOYAEhLjQM= github.com/golang/snappy v0.0.4/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q= @@ -155,6 +174,8 @@ github.com/google/btree v1.0.0/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ github.com/google/flatbuffers v2.0.0+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8= github.com/google/flatbuffers v22.10.26+incompatible h1:z1QiaMyPu1x3Z6xf2u1dsLj1ZxicdGSeaLpCuIsQNZM= github.com/google/flatbuffers v22.10.26+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8= +github.com/google/flatbuffers v23.5.26+incompatible h1:M9dgRyhJemaM4Sw8+66GHBu8ioaQmyPLg1b8VwK5WJg= +github.com/google/flatbuffers v23.5.26+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8= github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M= github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= @@ -167,15 +188,22 @@ github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38= github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/martian/v3 v3.2.1 h1:d8MncMlErDFTwQGBK1xhv026j9kqhvw1Qv9IbWT1VLQ= +github.com/google/martian/v3 v3.3.2 h1:IqNFLAmvJOgVlpdEBiQbDc2EwKW77amAycfTuWKdfvw= github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI= github.com/google/uuid v1.0.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I= github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/google/uuid v1.3.1 h1:KjJaJ9iWZ3jOFZIf1Lqf4laDRCasjl0BCmnEGxkdLb4= +github.com/google/uuid v1.3.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/googleapis/enterprise-certificate-proxy v0.2.0 h1:y8Yozv7SZtlU//QXbezB6QkpuE6jMD2/gfzk4AftXjs= github.com/googleapis/enterprise-certificate-proxy v0.2.0/go.mod h1:8C0jb7/mgJe/9KK8Lm7X9ctZC2t60YyIpYEI16jx0Qg= +github.com/googleapis/enterprise-certificate-proxy v0.2.3 h1:yk9/cqRKtT9wXZSsRH9aurXEpJX+U6FLtpYTdC3R06k= +github.com/googleapis/enterprise-certificate-proxy v0.2.3/go.mod h1:AwSRAtLfXpU5Nm3pW+v7rGDHp09LsPtGY9MduiEsR9k= github.com/googleapis/gax-go/v2 v2.7.0 h1:IcsPKeInNvYi7eqSaDjiZqDDKu5rsmunY0Y1YupQSSQ= github.com/googleapis/gax-go/v2 v2.7.0/go.mod h1:TEop28CZZQ2y+c0VxMUmu1lV+fQx57QpBWsYpwqHJx8= +github.com/googleapis/gax-go/v2 v2.11.0 h1:9V9PWXEsWnPpQhu/PeQIkS4eGzMlTLGgt80cUUI8Ki4= +github.com/googleapis/gax-go/v2 v2.11.0/go.mod h1:DxmR61SGKkGLa2xigwuZIQpkCI2S5iydzRfb3peWZJI= github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY= github.com/gorilla/context v1.1.1/go.mod h1:kBGZzfjB9CEq2AlWe17Uuf7NDRt0dE0s8S51q0aT7Yg= github.com/gorilla/mux v1.6.2/go.mod h1:1lud6UwP+6orDFRuTfBEV8e9/aOM/c4fVVCaMa2zaAs= @@ -232,9 +260,13 @@ github.com/klauspost/asmfmt v1.3.2/go.mod h1:AG8TuvYojzulgDAMCnYn50l/5QV3Bs/tp6j github.com/klauspost/compress v1.13.6/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk= github.com/klauspost/compress v1.15.12 h1:YClS/PImqYbn+UILDnqxQCZ3RehC9N318SU3kElDUEM= github.com/klauspost/compress v1.15.12/go.mod h1:QPwzmACJjUTFsnSHH934V6woptycfrDDJnH7hvFVbGM= +github.com/klauspost/compress v1.16.7 h1:2mk3MPGNzKyxErAw8YaohYh69+pa4sIQSC0fPGCFR9I= +github.com/klauspost/compress v1.16.7/go.mod h1:ntbaceVETuRiXiv4DpjP66DpAtAGkEQskQzEyD//IeE= github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg= github.com/klauspost/cpuid/v2 v2.1.2 h1:XhdX4fqAJUA0yj+kUwMavO0hHrSPAecYdYf1ZmxHvak= github.com/klauspost/cpuid/v2 v2.1.2/go.mod h1:RVVoqg1df56z8g3pUjL/3lE5UfnlrJX8tyFgg4nqhuY= +github.com/klauspost/cpuid/v2 v2.2.5 h1:0E5MSMDEoAulmXNFquVs//DdoomxaoTY1kUhbc/qbZg= +github.com/klauspost/cpuid/v2 v2.2.5/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws= github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFBFZlji/RkVcI2GknAs/DXo4wKdlNEc= github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= @@ -310,6 +342,8 @@ github.com/pierrec/lz4 v2.0.5+incompatible/go.mod h1:pdkljMzZIN41W+lC3N2tnIh5sFi github.com/pierrec/lz4/v4 v4.1.12/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4= github.com/pierrec/lz4/v4 v4.1.17 h1:kV4Ip+/hUBC+8T6+2EgburRtkE9ef4nbY3f4dFhGjMc= github.com/pierrec/lz4/v4 v4.1.17/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4= +github.com/pierrec/lz4/v4 v4.1.18 h1:xaKrnTkyoqfh1YItXl56+6KJNVYWlEEPuAQW9xsplYQ= +github.com/pierrec/lz4/v4 v4.1.18/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4= github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= @@ -377,6 +411,8 @@ github.com/stretchr/testify v1.7.2/go.mod h1:R6va5+xMeoiuVRoj+gSkQ7d3FALtqAAGI1F github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk= github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= +github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk= +github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= github.com/thoas/go-funk v0.9.1 h1:O549iLZqPpTUQ10ykd26sZhzD+rmR5pWhuElrhbC20M= github.com/tmc/grpc-websocket-proxy v0.0.0-20170815181823-89b8d40f7ca8/go.mod h1:ncp9v5uamzpCO7NfCPTXjqaC+bZgJeR0sMTm6dMHP7U= github.com/urfave/cli v1.20.0/go.mod h1:70zkFmudgCuE/ngEzBv17Jvp/497gISqfk5gWijbERA= @@ -435,6 +471,8 @@ golang.org/x/exp v0.0.0-20191002040644-a1355ae1e2c3/go.mod h1:NOZ3BPKG0ec/BKJQgn golang.org/x/exp v0.0.0-20211216164055-b2b84827b756/go.mod h1:b9TAUYHmRtqA6klRHApnXMnj+OyLce4yF5cZCUbk2ps= golang.org/x/exp v0.0.0-20221031165847-c99f073a8326 h1:QfTh0HpN6hlw6D3vu8DAwC8pBIwikq0AI1evdm+FksE= golang.org/x/exp v0.0.0-20221031165847-c99f073a8326/go.mod h1:CxIveKay+FTh1D0yPZemJVgC/95VzuuOLq5Qi4xnoYc= +golang.org/x/exp v0.0.0-20231006140011-7918f672742d h1:jtJma62tbqLibJ5sFQz8bKtEM8rJBtfilJ2qTU199MI= +golang.org/x/exp v0.0.0-20231006140011-7918f672742d/go.mod h1:ldy0pHrwJyGW56pPQzzkH36rKxoZW1tw7ZJpeKx+hdo= golang.org/x/image v0.0.0-20180708004352-c73c2afc3b81/go.mod h1:ux5Hcp/YLpHSI86hEcLt0YII63i6oz57MZXIpbrjZUs= golang.org/x/image v0.0.0-20190227222117-0694c2d4d067/go.mod h1:kZ7UVZpmo3dzQBMxlp+ypCbDeSB+sBbTgSJuh5dn5js= golang.org/x/image v0.0.0-20190802002840-cff245a6509b/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0= @@ -465,6 +503,8 @@ golang.org/x/mod v0.6.0-dev.0.20220106191415-9b9b3d81d5e3/go.mod h1:3p9vT2HGsQu2 golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4= golang.org/x/mod v0.6.0 h1:b9gGHsz9/HhJ3HF5DHQytPpuwocVTChQJK3AvoLRD5I= golang.org/x/mod v0.6.0/go.mod h1:4mET923SAdbXp2ki8ey+zGs1SLqsuM2Y0uvdZR/fUNI= +golang.org/x/mod v0.13.0 h1:I/DsJXRlw/8l/0c24sM9yb0T4z9liZTduXvdAWYiysY= +golang.org/x/mod v0.13.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c= golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= @@ -491,11 +531,15 @@ golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd/go.mod h1:CfG3xpIq0wQ8r1q4Su golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c= golang.org/x/net v0.1.0 h1:hZ/3BUoy5aId7sCpA/Tc5lt8DkFgdVS2onTpJsZ/fl0= golang.org/x/net v0.1.0/go.mod h1:Cx3nUiGt4eDBEyega/BKRp+/AlGL8hYe7U9odMt2Cco= +golang.org/x/net v0.17.0 h1:pVaXccu2ozPjCXewfr1S7xza/zcXTity9cCdXQYSjIM= +golang.org/x/net v0.17.0/go.mod h1:NxSsAGuq816PNPmqtQdLE42eU2Fs7NoRIZrHJAlaCOE= golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw= golang.org/x/oauth2 v0.0.0-20200107190931-bf48bf16ab8d/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw= golang.org/x/oauth2 v0.1.0 h1:isLCZuhj4v+tYv7eskaN4v/TM+A1begWWgyVJDdl1+Y= golang.org/x/oauth2 v0.1.0/go.mod h1:G9FE4dLTsbXUu90h/Pf85g4w1D+SSAgR+q46nJZ8M4A= +golang.org/x/oauth2 v0.10.0 h1:zHCpF2Khkwy4mMB4bv0U37YtJdTGW8jI0glAApi0Kh8= +golang.org/x/oauth2 v0.10.0/go.mod h1:kTpgurOux7LqtuxjuyZa4Gj2gdezIt/jQtGnNFfypQI= golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= @@ -539,6 +583,9 @@ golang.org/x/sys v0.0.0-20220704084225-05e143d24a9e/go.mod h1:oPkhp1MJrh7nUepCBc golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.1.0 h1:kunALQeHf1/185U1i0GOB/fy1IPRDDpuoOOqRReG57U= golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.13.0 h1:Af8nKPmuFypiUBjVoU9V20FiaFXOcuZI21p0ycVYYGE= +golang.org/x/sys v0.13.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= golang.org/x/term v0.1.0/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= @@ -550,6 +597,8 @@ golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= golang.org/x/text v0.4.0 h1:BrVqGRd7+k1DiOgtnFvAkoQEWQvBc25ouMJM6429SFg= golang.org/x/text v0.4.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= +golang.org/x/text v0.13.0 h1:ablQoSUd0tRdKxZewP80B+BaqeKJuVhuRxj/dkrun3k= +golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE= golang.org/x/time v0.0.0-20180412165947-fbb02b2291d2/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/tools v0.0.0-20180221164845-07fd8470d635/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= @@ -579,6 +628,8 @@ golang.org/x/tools v0.1.10/go.mod h1:Uh6Zz+xoGYZom868N8YTex3t7RhtHDBrE8Gzo9bV56E golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc= golang.org/x/tools v0.2.0 h1:G6AHpWxTMGY1KyEYoAQ5WTtIekUUvDNjan3ugu60JvE= golang.org/x/tools v0.2.0/go.mod h1:y4OqIKeOV/fWJetJ8bXPU1sEVniLMIyDAZWeHdV+NTA= +golang.org/x/tools v0.14.0 h1:jvNa2pY0M4r62jkRQ6RwEZZyPcymeL9XZMLBbV7U2nc= +golang.org/x/tools v0.14.0/go.mod h1:uYBEerGOWcJyEORxN+Ek8+TT266gXkNlHdJBwexUsBg= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= @@ -590,12 +641,15 @@ gonum.org/v1/gonum v0.0.0-20180816165407-929014505bf4/go.mod h1:Y+Yx5eoAFn32cQvJ gonum.org/v1/gonum v0.8.2/go.mod h1:oe/vMfY3deqTw+1EZJhuvEW2iwGF1bW9wwu7XCu0+v0= gonum.org/v1/gonum v0.9.3 h1:DnoIG+QAMaF5NvxnGe/oKsgKcAc6PcUyl8q0VetfQ8s= gonum.org/v1/gonum v0.9.3/go.mod h1:TZumC3NeyVQskjXqmyWt4S3bINhy7B4eYwW69EbyX+0= +gonum.org/v1/gonum v0.12.0 h1:xKuo6hzt+gMav00meVPUlXwSdoEJP46BR+wdxQEFK2o= gonum.org/v1/netlib v0.0.0-20190313105609-8cb42192e0e0/go.mod h1:wa6Ws7BG/ESfp6dHfk7C6KdzKA7wR7u/rKwOGE66zvw= gonum.org/v1/plot v0.0.0-20190515093506-e2840ee46a6b/go.mod h1:Wt8AAjI+ypCyYX3nZBvf6cAIx93T+c/OS2HFAYskSZc= gonum.org/v1/plot v0.9.0/go.mod h1:3Pcqqmp6RHvJI72kgb8fThyUnav364FOsdDo2aGW5lY= google.golang.org/api v0.3.1/go.mod h1:6wY9I6uQWHQ8EM57III9mq/AjF+i8G65rmVagqKMtkk= google.golang.org/api v0.102.0 h1:JxJl2qQ85fRMPNvlZY/enexbxpCjLwGhZUtgfGeQ51I= google.golang.org/api v0.102.0/go.mod h1:3VFl6/fzoA+qNuS1N1/VfXY4LjoXN/wzeIp7TweWwGo= +google.golang.org/api v0.126.0 h1:q4GJq+cAdMAC7XP7njvQ4tvohGLiSlytuL4BQxbIZ+o= +google.golang.org/api v0.126.0/go.mod h1:mBwVAtz+87bEN6CbA1GtZPDOqY2R5ONPqJeIlvyo4Aw= google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM= google.golang.org/appengine v1.2.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4= google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4= @@ -611,6 +665,8 @@ google.golang.org/genproto v0.0.0-20200526211855-cb27e3aa2013/go.mod h1:NbSheEEY google.golang.org/genproto v0.0.0-20210630183607-d20f26d13c79/go.mod h1:yiaVoXHpRzHGyxV3o4DktVWY4mSUErTKaeEOq6C3t3U= google.golang.org/genproto v0.0.0-20221027153422-115e99e71e1c h1:QgY/XxIAIeccR+Ca/rDdKubLIU9rcJ3xfy1DC/Wd2Oo= google.golang.org/genproto v0.0.0-20221027153422-115e99e71e1c/go.mod h1:CGI5F/G+E5bKwmfYo09AXuVN4dD894kIKUFmVbP2/Fo= +google.golang.org/genproto v0.0.0-20230711160842-782d3b101e98 h1:Z0hjGZePRE0ZBWotvtrwxFNrNE9CUAGtplaDK5NNI/g= +google.golang.org/genproto v0.0.0-20230711160842-782d3b101e98/go.mod h1:S7mY02OqCJTD0E1OiQy1F72PWFB4bZJ87cAtLPYgDR0= google.golang.org/grpc v1.17.0/go.mod h1:6QZJwpn2B+Zp71q/5VxRsJ6NXXVCE5NRUHRo+f3cWCs= google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c= google.golang.org/grpc v1.20.0/go.mod h1:chYK+tFQF0nDUGJgXMSgLCQk3phJEuONr2DCgLDdAQM= @@ -629,6 +685,8 @@ google.golang.org/grpc v1.38.0/go.mod h1:NREThFqKR1f3iQ6oBuvc5LadQuXVGo9rkm5ZGrQ google.golang.org/grpc v1.41.0/go.mod h1:U3l9uK9J0sini8mHphKoXyaqDA/8VyGnDee1zzIUK6k= google.golang.org/grpc v1.50.1 h1:DS/BukOZWp8s6p4Dt/tOaJaTQyPyOoCcrjroHuCeLzY= google.golang.org/grpc v1.50.1/go.mod h1:ZgQEeidpAuNRZ8iRrlBKXZQP1ghovWIVhdJRyCDK+GI= +google.golang.org/grpc v1.58.2 h1:SXUpjxeVF3FKrTYQI4f4KvbGD5u2xccdYdurwowix5I= +google.golang.org/grpc v1.58.2/go.mod h1:tgX3ZQDlNJGU96V6yHh1T/JeoBQ2TXdr43YbYSsCJk0= google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8= google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0= google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM= @@ -643,6 +701,8 @@ google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQ google.golang.org/protobuf v1.27.1/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc= google.golang.org/protobuf v1.28.1 h1:d0NfwRgPtno5B1Wa6L2DAG+KivqkdutMf1UhdNx175w= google.golang.org/protobuf v1.28.1/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= +google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs8= +google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=