@@ -1209,47 +1209,6 @@ class IntervalColumnWriter : public BasicColumnWriter {
12091209 }
12101210};
12111211
1212- // ===--------------------------------------------------------------------===//
1213- // Geometry Column Writer
1214- // ===--------------------------------------------------------------------===//
1215- // This class just wraps another column writer, but also calculates the extent
1216- // of the geometry column by updating the geodata object with every written
1217- // vector.
1218- template <class WRITER_IMPL >
1219- class GeometryColumnWriter : public WRITER_IMPL {
1220- GeoParquetColumnMetadata geo_data;
1221- GeoParquetColumnMetadataWriter geo_data_writer;
1222- string column_name;
1223-
1224- public:
1225- void Write (ColumnWriterState &state, Vector &vector, idx_t count) override {
1226- // Just write normally
1227- WRITER_IMPL::Write (state, vector, count);
1228-
1229- // And update the geodata object
1230- geo_data_writer.Update (geo_data, vector, count);
1231- }
1232- void FinalizeWrite (ColumnWriterState &state) override {
1233- WRITER_IMPL::FinalizeWrite (state);
1234-
1235- // Add the geodata object to the writer
1236- this ->writer .GetGeoParquetData ().geometry_columns [column_name] = geo_data;
1237- }
1238-
1239- public:
1240- GeometryColumnWriter (ClientContext &context, ParquetWriter &writer, idx_t schema_idx, vector<string> schema_path_p,
1241- idx_t max_repeat, idx_t max_define, bool can_have_nulls, string name)
1242- : WRITER_IMPL(writer, schema_idx, std::move(schema_path_p), max_repeat, max_define, can_have_nulls),
1243- geo_data_writer (context), column_name(std::move(name)) {
1244-
1245- auto &geo_data = writer.GetGeoParquetData ();
1246- if (geo_data.primary_geometry_column .empty ()) {
1247- // Set the first column to the primary column
1248- geo_data.primary_geometry_column = column_name;
1249- }
1250- }
1251- };
1252-
12531212// ===--------------------------------------------------------------------===//
12541213// String Column Writer
12551214// ===--------------------------------------------------------------------===//
@@ -1563,6 +1522,58 @@ class StringColumnWriter : public BasicColumnWriter {
15631522 }
15641523};
15651524
1525+ // ===--------------------------------------------------------------------===//
1526+ // WKB Column Writer
1527+ // ===--------------------------------------------------------------------===//
1528+ // Used to store the metadata for a WKB-encoded geometry column when writing
1529+ // GeoParquet files.
1530+ class WKBColumnWriterState final : public StringColumnWriterState {
1531+ public:
1532+ WKBColumnWriterState (ClientContext &context, duckdb_parquet::format::RowGroup &row_group, idx_t col_idx)
1533+ : StringColumnWriterState(row_group, col_idx), geo_data(), geo_data_writer(context) {
1534+ }
1535+
1536+ GeoParquetColumnMetadata geo_data;
1537+ GeoParquetColumnMetadataWriter geo_data_writer;
1538+ };
1539+
1540+ class WKBColumnWriter final : public StringColumnWriter {
1541+ public:
1542+ WKBColumnWriter (ClientContext &context_p, ParquetWriter &writer, idx_t schema_idx, vector<string> schema_path_p,
1543+ idx_t max_repeat, idx_t max_define, bool can_have_nulls, string name)
1544+ : StringColumnWriter(writer, schema_idx, std::move(schema_path_p), max_repeat, max_define, can_have_nulls),
1545+ column_name (std::move(name)), context(context_p) {
1546+
1547+ this ->writer .GetGeoParquetData ().RegisterGeometryColumn (column_name);
1548+ }
1549+
1550+ unique_ptr<ColumnWriterState> InitializeWriteState (duckdb_parquet::format::RowGroup &row_group) override {
1551+ auto result = make_uniq<WKBColumnWriterState>(context, row_group, row_group.columns .size ());
1552+ RegisterToRowGroup (row_group);
1553+ return std::move (result);
1554+ }
1555+ void Write (ColumnWriterState &state, Vector &vector, idx_t count) override {
1556+ StringColumnWriter::Write (state, vector, count);
1557+
1558+ auto &geo_state = state.Cast <WKBColumnWriterState>();
1559+ geo_state.geo_data_writer .Update (geo_state.geo_data , vector, count);
1560+ }
1561+
1562+ void FinalizeWrite (ColumnWriterState &state) override {
1563+ StringColumnWriter::FinalizeWrite (state);
1564+
1565+ // Add the geodata object to the writer
1566+ const auto &geo_state = state.Cast <WKBColumnWriterState>();
1567+
1568+ // Merge this state's geo column data with the writer's geo column data
1569+ writer.GetGeoParquetData ().FlushColumnMeta (column_name, geo_state.geo_data );
1570+ }
1571+
1572+ private:
1573+ string column_name;
1574+ ClientContext &context;
1575+ };
1576+
15661577// ===--------------------------------------------------------------------===//
15671578// Enum Column Writer
15681579// ===--------------------------------------------------------------------===//
@@ -2234,8 +2245,8 @@ unique_ptr<ColumnWriter> ColumnWriter::CreateWriterRecursive(ClientContext &cont
22342245 schema_path.push_back (name);
22352246
22362247 if (type.id () == LogicalTypeId::BLOB && type.GetAlias () == " WKB_BLOB"
2237- return make_uniq<GeometryColumnWriter<StringColumnWriter>> (context, writer, schema_idx, std::move (schema_path),
2238- max_repeat, max_define, can_have_nulls, name);
2248+ return make_uniq<WKBColumnWriter> (context, writer, schema_idx, std::move (schema_path), max_repeat, max_define ,
2249+ can_have_nulls, name);
22392250 }
22402251
22412252 switch (type.id ()) {
0 commit comments