Skip to content

Commit

Permalink
use array writer in kernel
Browse files Browse the repository at this point in the history
  • Loading branch information
@paleolimbot
paleolimbot committed Sep 17, 2023
1 parent c333a5f commit c9c4c8b
Show file tree
Hide file tree
Showing 2 changed files with 53 additions and 81 deletions.
106 changes: 29 additions & 77 deletions src/geoarrow/kernel.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -101,9 +101,8 @@ struct GeoArrowVisitorKernelPrivate {
int visit_by_feature;
struct GeoArrowArrayReader reader;
struct GeoArrowArrayView array_view;
struct GeoArrowWKBWriter wkb_writer;
struct GeoArrowArrayWriter writer;
struct GeoArrowWKTWriter wkt_writer;
struct GeoArrowBuilder builder;
struct GeoArrowGeometryTypesVisitorPrivate geometry_types_private;
struct GeoArrowBox2DPrivate box2d_private;
int (*finish_push_batch)(struct GeoArrowVisitorKernelPrivate* private_data,
Expand Down Expand Up @@ -146,12 +145,8 @@ static void kernel_release_visitor(struct GeoArrowKernel* kernel) {
GeoArrowArrayReaderReset(&private_data->reader);
}

if (private_data->builder.private_data != NULL) {
GeoArrowBuilderReset(&private_data->builder);
}

if (private_data->wkb_writer.private_data != NULL) {
GeoArrowWKBWriterReset(&private_data->wkb_writer);
if (private_data->writer.private_data != NULL) {
GeoArrowArrayWriterReset(&private_data->writer);
}

if (private_data->wkt_writer.private_data != NULL) {
Expand Down Expand Up @@ -251,24 +246,10 @@ static int finish_start_visit_void_agg(struct GeoArrowVisitorKernelPrivate* priv
return ArrowSchemaInitFromType(out, NANOARROW_TYPE_NA);
}

// Kernels as_wkt and format_wkt
// Kernel format_wkt
//
// Visits every feature in the input and writes the corresponding well-known text output.
// For the format_wkt kernel, optionally specify precision and
// max_element_size_bytes.

static int finish_start_as_wkt(struct GeoArrowVisitorKernelPrivate* private_data,
struct ArrowSchema* schema, const char* options,
struct ArrowSchema* out, struct GeoArrowError* error) {
GeoArrowWKTWriterInitVisitor(&private_data->wkt_writer, &private_data->v);

struct ArrowSchema tmp;
NANOARROW_RETURN_NOT_OK(GeoArrowSchemaInitExtension(&tmp, GEOARROW_TYPE_WKT));
NANOARROW_RETURN_NOT_OK(GeoArrowSchemaSetMetadataFrom(&tmp, schema));
ArrowSchemaMove(&tmp, out);

return GEOARROW_OK;
}
// Visits every feature in the input and writes the corresponding well-known text output,
// optionally specifying precision and max_element_size_bytes.

static int finish_start_format_wkt(struct GeoArrowVisitorKernelPrivate* private_data,
struct ArrowSchema* schema, const char* options,
Expand All @@ -285,41 +266,19 @@ static int finish_start_format_wkt(struct GeoArrowVisitorKernelPrivate* private_

GeoArrowWKTWriterInitVisitor(&private_data->wkt_writer, &private_data->v);

struct ArrowSchema tmp;
NANOARROW_RETURN_NOT_OK(ArrowSchemaInitFromType(&tmp, NANOARROW_TYPE_STRING));
ArrowSchemaMove(&tmp, out);

NANOARROW_RETURN_NOT_OK(ArrowSchemaInitFromType(out, NANOARROW_TYPE_STRING));
return GEOARROW_OK;
}

static int finish_push_batch_as_wkt(struct GeoArrowVisitorKernelPrivate* private_data,
struct ArrowArray* out, struct GeoArrowError* error) {
static int finish_push_batch_format_wkt(struct GeoArrowVisitorKernelPrivate* private_data,
struct ArrowArray* out,
struct GeoArrowError* error) {
return GeoArrowWKTWriterFinish(&private_data->wkt_writer, out, error);
}

// Kernel as_wkb
//
// Visits every feature in the input and writes the corresponding well-known binary output

static int finish_start_as_wkb(struct GeoArrowVisitorKernelPrivate* private_data,
struct ArrowSchema* schema, const char* options,
struct ArrowSchema* out, struct GeoArrowError* error) {
struct ArrowSchema tmp;
NANOARROW_RETURN_NOT_OK(GeoArrowSchemaInitExtension(&tmp, GEOARROW_TYPE_WKB));
NANOARROW_RETURN_NOT_OK(GeoArrowSchemaSetMetadataFrom(&tmp, schema));
ArrowSchemaMove(&tmp, out);

return GEOARROW_OK;
}

static int finish_push_batch_as_wkb(struct GeoArrowVisitorKernelPrivate* private_data,
struct ArrowArray* out, struct GeoArrowError* error) {
return GeoArrowWKBWriterFinish(&private_data->wkb_writer, out, error);
}

// Kernel as_geoarrow
//
// Visits every feature in the input and writes a geoarrow-encoded array.
// Visits every feature in the input and writes an array of the specified type.
// Takes option 'type' as the desired integer enum GeoArrowType.

static int finish_start_as_geoarrow(struct GeoArrowVisitorKernelPrivate* private_data,
Expand All @@ -330,26 +289,34 @@ static int finish_start_as_geoarrow(struct GeoArrowVisitorKernelPrivate* private
NANOARROW_RETURN_NOT_OK(kernel_get_arg_long(options, "type", &out_type_long, 1, error));
enum GeoArrowType out_type = (enum GeoArrowType)out_type_long;

if (out_type != private_data->builder.view.schema_view.type) {
GeoArrowBuilderReset(&private_data->builder);
NANOARROW_RETURN_NOT_OK(
GeoArrowBuilderInitFromType(&private_data->builder, out_type));
NANOARROW_RETURN_NOT_OK(
GeoArrowBuilderInitVisitor(&private_data->builder, &private_data->v));
if (private_data->writer.private_data != NULL) {
GeoArrowErrorSet(error, "Expected exactly one call to start(as_geoarrow)");
return EINVAL;
}

NANOARROW_RETURN_NOT_OK(
GeoArrowArrayWriterInitFromType(&private_data->writer, out_type));
NANOARROW_RETURN_NOT_OK(
GeoArrowArrayWriterInitVisitor(&private_data->writer, &private_data->v));

struct ArrowSchema tmp;
NANOARROW_RETURN_NOT_OK(GeoArrowSchemaInitExtension(&tmp, out_type));
NANOARROW_RETURN_NOT_OK(GeoArrowSchemaSetMetadataFrom(&tmp, schema));
ArrowSchemaMove(&tmp, out);

int result = GeoArrowSchemaSetMetadataFrom(&tmp, schema);
if (result != GEOARROW_OK) {
GeoArrowErrorSet(error, "GeoArrowSchemaSetMetadataFrom() failed");
tmp.release(&tmp);
return result;
}

ArrowSchemaMove(&tmp, out);
return GEOARROW_OK;
}

static int finish_push_batch_as_geoarrow(
struct GeoArrowVisitorKernelPrivate* private_data, struct ArrowArray* out,
struct GeoArrowError* error) {
return GeoArrowBuilderFinish(&private_data->builder, out, error);
return GeoArrowArrayWriterFinish(&private_data->writer, out, error);
}

// Kernel unique_geometry_types_agg
Expand Down Expand Up @@ -694,23 +661,12 @@ static int GeoArrowInitVisitorKernelInternal(struct GeoArrowKernel* kernel,
if (strcmp(name, "visit_void_agg") == 0) {
kernel->finish = &kernel_finish_void_agg;
private_data->finish_start = &finish_start_visit_void_agg;
} else if (strcmp(name, "as_wkt") == 0) {
kernel->finish = &kernel_finish_void;
private_data->finish_start = &finish_start_as_wkt;
private_data->finish_push_batch = &finish_push_batch_as_wkt;
result = GeoArrowWKTWriterInit(&private_data->wkt_writer);
} else if (strcmp(name, "format_wkt") == 0) {
kernel->finish = &kernel_finish_void;
private_data->finish_start = &finish_start_format_wkt;
private_data->finish_push_batch = &finish_push_batch_as_wkt;
private_data->finish_push_batch = &finish_push_batch_format_wkt;
result = GeoArrowWKTWriterInit(&private_data->wkt_writer);
private_data->visit_by_feature = 1;
} else if (strcmp(name, "as_wkb") == 0) {
kernel->finish = &kernel_finish_void;
private_data->finish_start = &finish_start_as_wkb;
private_data->finish_push_batch = &finish_push_batch_as_wkb;
result = GeoArrowWKBWriterInit(&private_data->wkb_writer);
GeoArrowWKBWriterInitVisitor(&private_data->wkb_writer, &private_data->v);
} else if (strcmp(name, "as_geoarrow") == 0) {
kernel->finish = &kernel_finish_void;
private_data->finish_start = &finish_start_as_geoarrow;
Expand Down Expand Up @@ -751,12 +707,8 @@ GeoArrowErrorCode GeoArrowKernelInit(struct GeoArrowKernel* kernel, const char*
return NANOARROW_OK;
} else if (strcmp(name, "visit_void_agg") == 0) {
return GeoArrowInitVisitorKernelInternal(kernel, name);
} else if (strcmp(name, "as_wkt") == 0) {
return GeoArrowInitVisitorKernelInternal(kernel, name);
} else if (strcmp(name, "format_wkt") == 0) {
return GeoArrowInitVisitorKernelInternal(kernel, name);
} else if (strcmp(name, "as_wkb") == 0) {
return GeoArrowInitVisitorKernelInternal(kernel, name);
} else if (strcmp(name, "as_geoarrow") == 0) {
return GeoArrowInitVisitorKernelInternal(kernel, name);
} else if (strcmp(name, "unique_geometry_types_agg") == 0) {
Expand Down
28 changes: 24 additions & 4 deletions src/geoarrow/kernel_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -204,8 +204,17 @@ TEST(KernelTest, KernelTestAsWKT) {
ASSERT_EQ(ArrowArrayAppendNull(&array_in, 1), GEOARROW_OK);
ASSERT_EQ(ArrowArrayFinishBuildingDefault(&array_in, nullptr), GEOARROW_OK);

EXPECT_EQ(GeoArrowKernelInit(&kernel, "as_wkt", nullptr), GEOARROW_OK);
EXPECT_EQ(kernel.start(&kernel, &schema_in, nullptr, &schema_out, &error), GEOARROW_OK);
EXPECT_EQ(GeoArrowKernelInit(&kernel, "as_geoarrow", nullptr), GEOARROW_OK);

struct ArrowBuffer buffer;
ASSERT_EQ(ArrowMetadataBuilderInit(&buffer, nullptr), GEOARROW_OK);
ASSERT_EQ(
ArrowMetadataBuilderAppend(&buffer, ArrowCharView("type"), ArrowCharView("100003")),
GEOARROW_OK);

EXPECT_EQ(kernel.start(&kernel, &schema_in, reinterpret_cast<char*>(buffer.data),
&schema_out, &error),
GEOARROW_OK);
EXPECT_STREQ(schema_out.format, "u");
EXPECT_EQ(kernel.push_batch(&kernel, &array_in, &array_out1, &error), GEOARROW_OK);
EXPECT_EQ(kernel.push_batch(&kernel, &array_in, &array_out2, &error), GEOARROW_OK);
Expand All @@ -231,6 +240,7 @@ TEST(KernelTest, KernelTestAsWKT) {
EXPECT_EQ(std::string(item.data, item.size_bytes), "POINT (30 10)");
EXPECT_TRUE(ArrowArrayViewIsNull(&array_view, 1));

ArrowBufferReset(&buffer);
ArrowArrayViewReset(&array_view);
schema_in.release(&schema_in);
schema_out.release(&schema_out);
Expand Down Expand Up @@ -446,8 +456,17 @@ TEST(KernelTest, KernelTestAsWKB) {
ASSERT_EQ(ArrowArrayAppendNull(&array_in, 1), GEOARROW_OK);
ASSERT_EQ(ArrowArrayFinishBuildingDefault(&array_in, nullptr), GEOARROW_OK);

EXPECT_EQ(GeoArrowKernelInit(&kernel, "as_wkb", nullptr), GEOARROW_OK);
EXPECT_EQ(kernel.start(&kernel, &schema_in, nullptr, &schema_out, &error), GEOARROW_OK);
EXPECT_EQ(GeoArrowKernelInit(&kernel, "as_geoarrow", nullptr), GEOARROW_OK);

struct ArrowBuffer buffer;
ASSERT_EQ(ArrowMetadataBuilderInit(&buffer, nullptr), GEOARROW_OK);
ASSERT_EQ(
ArrowMetadataBuilderAppend(&buffer, ArrowCharView("type"), ArrowCharView("100001")),
GEOARROW_OK);

EXPECT_EQ(kernel.start(&kernel, &schema_in, reinterpret_cast<char*>(buffer.data),
&schema_out, &error),
GEOARROW_OK);
EXPECT_STREQ(schema_out.format, "z");
EXPECT_EQ(kernel.push_batch(&kernel, &array_in, &array_out1, &error), GEOARROW_OK);
EXPECT_EQ(kernel.push_batch(&kernel, &array_in, &array_out2, &error), GEOARROW_OK);
Expand Down Expand Up @@ -480,6 +499,7 @@ TEST(KernelTest, KernelTestAsWKB) {
EXPECT_EQ(std::basic_string<uint8_t>(item.data.as_uint8, item.size_bytes), point);
EXPECT_TRUE(ArrowArrayViewIsNull(&array_view, 1));

ArrowBufferReset(&buffer);
ArrowArrayViewReset(&array_view);
schema_in.release(&schema_in);
schema_out.release(&schema_out);
Expand Down

0 comments on commit c9c4c8b

Please sign in to comment.