Skip to content

Commit d381b52

Browse files
[Spark] fix spark read sec/milli timestamp arrow vector (lakesoul-io#576)
Signed-off-by: chenxu <chenxu@dmetasoul.com> Co-authored-by: chenxu <chenxu@dmetasoul.com>
1 parent 9df97d6 commit d381b52

File tree

1 file changed

+74
-6
lines changed

1 file changed

+74
-6
lines changed

lakesoul-spark/src/main/java/org/apache/spark/sql/arrow/ArrowColumnVector.java

Lines changed: 74 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -158,9 +158,17 @@ void initAccessor(ValueVector vector) {
158158
} else if (vector instanceof DateDayVector) {
159159
accessor = new DateAccessor((DateDayVector) vector);
160160
} else if (vector instanceof TimeStampMicroTZVector) {
161-
accessor = new TimestampAccessor((TimeStampMicroTZVector) vector);
161+
accessor = new TimestampMicroAccessor((TimeStampMicroTZVector) vector);
162162
} else if (vector instanceof TimeStampMicroVector) {
163-
accessor = new TimestampNTZAccessor((TimeStampMicroVector) vector);
163+
accessor = new TimestampMicroNTZAccessor((TimeStampMicroVector) vector);
164+
} else if (vector instanceof TimeStampMilliTZVector) {
165+
accessor = new TimestampMilliAccessor((TimeStampMilliTZVector) vector);
166+
} else if (vector instanceof TimeStampMilliVector) {
167+
accessor = new TimestampMilliNTZAccessor((TimeStampMilliVector) vector);
168+
} else if (vector instanceof TimeStampSecTZVector) {
169+
accessor = new TimestampSecAccessor((TimeStampSecTZVector) vector);
170+
} else if (vector instanceof TimeStampSecVector) {
171+
accessor = new TimestampSecNTZAccessor((TimeStampSecVector) vector);
164172
} else if (vector instanceof MapVector) {
165173
MapVector mapVector = (MapVector) vector;
166174
accessor = new MapAccessor(mapVector);
@@ -432,11 +440,11 @@ final int getInt(int rowId) {
432440
}
433441
}
434442

435-
static class TimestampAccessor extends ArrowVectorAccessor {
443+
static class TimestampMicroAccessor extends ArrowVectorAccessor {
436444

437445
private final TimeStampMicroTZVector accessor;
438446

439-
TimestampAccessor(TimeStampMicroTZVector vector) {
447+
TimestampMicroAccessor(TimeStampMicroTZVector vector) {
440448
super(vector);
441449
this.accessor = vector;
442450
}
@@ -447,11 +455,11 @@ final long getLong(int rowId) {
447455
}
448456
}
449457

450-
static class TimestampNTZAccessor extends ArrowVectorAccessor {
458+
static class TimestampMicroNTZAccessor extends ArrowVectorAccessor {
451459

452460
private final TimeStampMicroVector accessor;
453461

454-
TimestampNTZAccessor(TimeStampMicroVector vector) {
462+
TimestampMicroNTZAccessor(TimeStampMicroVector vector) {
455463
super(vector);
456464
this.accessor = vector;
457465
}
@@ -462,6 +470,66 @@ final long getLong(int rowId) {
462470
}
463471
}
464472

473+
static class TimestampSecAccessor extends ArrowVectorAccessor {
474+
475+
private final TimeStampSecTZVector accessor;
476+
477+
TimestampSecAccessor(TimeStampSecTZVector vector) {
478+
super(vector);
479+
this.accessor = vector;
480+
}
481+
482+
@Override
483+
final long getLong(int rowId) {
484+
return accessor.get(rowId) * 1000000;
485+
}
486+
}
487+
488+
static class TimestampSecNTZAccessor extends ArrowVectorAccessor {
489+
490+
private final TimeStampSecVector accessor;
491+
492+
TimestampSecNTZAccessor(TimeStampSecVector vector) {
493+
super(vector);
494+
this.accessor = vector;
495+
}
496+
497+
@Override
498+
final long getLong(int rowId) {
499+
return accessor.get(rowId) * 1000000;
500+
}
501+
}
502+
503+
static class TimestampMilliAccessor extends ArrowVectorAccessor {
504+
505+
private final TimeStampMilliTZVector accessor;
506+
507+
TimestampMilliAccessor(TimeStampMilliTZVector vector) {
508+
super(vector);
509+
this.accessor = vector;
510+
}
511+
512+
@Override
513+
final long getLong(int rowId) {
514+
return accessor.get(rowId) * 1000;
515+
}
516+
}
517+
518+
static class TimestampMilliNTZAccessor extends ArrowVectorAccessor {
519+
520+
private final TimeStampMilliVector accessor;
521+
522+
TimestampMilliNTZAccessor(TimeStampMilliVector vector) {
523+
super(vector);
524+
this.accessor = vector;
525+
}
526+
527+
@Override
528+
final long getLong(int rowId) {
529+
return accessor.get(rowId) * 1000;
530+
}
531+
}
532+
465533
static class ArrayAccessor extends ArrowVectorAccessor {
466534

467535
private final ListVector accessor;

0 commit comments

Comments
 (0)