Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(java): Move schema caching to unsafe trait to avoid issues when using non-inferred schema. #1944

Merged
merged 6 commits into from
Nov 19, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@ public BinaryArray(Field field) {
} else {
this.elementSize = width;
}
initializeExtData(1); // Only require at most one slot to cache the schema for array type.
}

public void pointTo(MemoryBuffer buffer, int offset, int sizeInBytes) {
Expand Down Expand Up @@ -135,7 +136,7 @@ public BigDecimal getDecimal(int ordinal) {

@Override
public BinaryRow getStruct(int ordinal) {
return getStruct(ordinal, field.getChildren().get(0));
return getStruct(ordinal, field.getChildren().get(0), 0);
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@ public BinaryRow(Schema schema) {
this.numFields = schema.getFields().size();
Preconditions.checkArgument(numFields > 0);
this.bitmapWidthInBytes = BitUtils.calculateBitmapWidthInBytes(numFields);
initializeExtData(numFields);
}

public void pointTo(MemoryBuffer buffer, int offset, int sizeInBytes) {
Expand Down Expand Up @@ -155,7 +156,7 @@ public BigDecimal getDecimal(int ordinal) {

@Override
public BinaryRow getStruct(int ordinal) {
return getStruct(ordinal, schema.getFields().get(ordinal));
return getStruct(ordinal, schema.getFields().get(ordinal), ordinal);
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.arrow.vector.util.DecimalUtility;
import org.apache.fury.format.row.Getters;
import org.apache.fury.format.row.Setters;
Expand All @@ -35,6 +36,7 @@

/** Internal to binary row format to reuse code, don't use it in anywhere else. */
abstract class UnsafeTrait implements Getters, Setters {
private Object[] extData;

abstract MemoryBuffer getBuffer();

Expand All @@ -55,6 +57,10 @@ public MemoryBuffer getBuffer(int ordinal) {

abstract int getOffset(int ordinal);

void initializeExtData(int numSlots) {
extData = new Object[numSlots];
}

// ###########################################################
// ####################### getters #######################
// ###########################################################
Expand Down Expand Up @@ -143,14 +149,25 @@ BigDecimal getDecimal(int ordinal, ArrowType.Decimal decimalType) {
return decimal;
}

BinaryRow getStruct(int ordinal, Field field) {
/**
* Gets the field at a specific ordinal as a struct.
*
* @param ordinal the ordinal position of this field.
* @param field the Arrow field corresponding to this struct.
* @param extDataSlot the ext data slot used to cache the schema for the struct.
* @return the binary row representation of the struct.
*/
BinaryRow getStruct(int ordinal, Field field, int extDataSlot) {
if (isNullAt(ordinal)) {
return null;
}
final long offsetAndSize = getInt64(ordinal);
final int relativeOffset = (int) (offsetAndSize >> 32);
final int size = (int) offsetAndSize;
BinaryRow row = new BinaryRow(DataTypes.createSchema(field));
if (extData[extDataSlot] == null) {
extData[extDataSlot] = DataTypes.createSchema(field);
}
BinaryRow row = new BinaryRow((Schema) extData[extDataSlot]);
row.pointTo(getBuffer(), getBaseOffset() + relativeOffset, size);
return row;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -314,7 +314,7 @@ public static Field field(String name, ArrowType type, Field... children) {
}

public static Field field(String name, FieldType fieldType, List<Field> children) {
return new ExtField(name, fieldType, children);
return new Field(name, fieldType, children);
}

public static Field notNullField(String name, ArrowType type, Field... children) {
Expand Down Expand Up @@ -396,19 +396,11 @@ public static Field mapField(String name, Field keyField, Field itemField) {
}

public static Field keyFieldForMap(Field mapField) {
Field field = mapField.getChildren().get(0).getChildren().get(0);
if (field.getClass() != ExtField.class) {
return new ExtField(field.getName(), field.getFieldType(), field.getChildren());
}
return field;
return mapField.getChildren().get(0).getChildren().get(0);
}

public static Field itemFieldForMap(Field mapField) {
Field field = mapField.getChildren().get(0).getChildren().get(1);
if (field.getClass() != ExtField.class) {
return new ExtField(field.getName(), field.getFieldType(), field.getChildren());
}
return field;
return mapField.getChildren().get(0).getChildren().get(1);
}

public static Field keyArrayFieldForMap(Field mapField) {
Expand All @@ -425,24 +417,7 @@ public static Schema schemaFromStructField(Field structField) {
}

public static Schema createSchema(Field field) {
if (field.getClass() != ExtField.class) {
throw new IllegalArgumentException(
String.format("Field %s got wrong type %s", field, field.getClass()));
}
ExtField extField = (ExtField) field;
Object extData = extField.extData;
if (extData == null) {
extField.extData = extData = new Schema(field.getChildren(), field.getMetadata());
}
return (Schema) extData;
}

static class ExtField extends Field {
Object extData;

public ExtField(String name, FieldType fieldType, List<Field> children) {
super(name, fieldType, children);
}
return new Schema(field.getChildren(), field.getMetadata());
}

public static Field structField(boolean nullable, Field... fields) {
Expand Down
Loading