From 777ba63c240f118f86c2b62570e117503935201e Mon Sep 17 00:00:00 2001 From: MisterRaindrop <278811821@qq.com> Date: Thu, 3 Sep 2026 15:15:21 +0800 Subject: [PATCH] datalake_fdw: read and write Parquet through Arrow The format layer had an interface and no implementation. This is the Parquet one, and the conversions either side of it: PostgreSQL tuples into an Arrow batch, and an Arrow batch back into Datums. Arrow rather than libparquet alone, because libparquet is written in terms of Arrow's types, so linking one links the other. It is a build dependency now, found with pkg-config -- and whatever -std= its .pc file asks for is filtered out, because those flags land after CXXFLAGS and the Arrow project's own packages say -std=c++11. A fragment is a range of row groups rather than a whole file, so several segments can read one large file. Reading is single-threaded: a worker thread that fails has no way to report it through PostgreSQL. A reader and a writer hold a descriptor and memory from Arrow's allocator, which transaction abort does not reclaim, so both register with the resource owner -- the shape PAX uses in comm/pax_resource.cc. Types: bool, the four integers, both floats, text, varchar, char, bytea, date, timestamp, timestamptz. Anything else is refused by name. The 30-year epoch shift is range checked in both directions: PostgreSQL's range runs past what Arrow holds as microseconds from 1970, and a round trip through two unchecked halves would agree with itself. datalake_fdw_test is a second extension in the same library with the two functions this can be run with from SQL. Its regression case round-trips every supported type and checks that reading row groups separately gives back what reading the whole file does. Arrow 9.0.0 (EPEL 9) builds and passes; 17.0.0 (Rocky 10, gcc 14) and 17.0.0 and 21.0.0 (Rocky 8, gcc 8) compile without warnings. What is written here reads back correctly in pyarrow 21, which is not the implementation that wrote it. --- .github/workflows/build-cloudberry.yml | 28 ++ contrib/datalake_fdw/Makefile | 82 +++- .../datalake_fdw/datalake_fdw_test--1.0.sql | 51 ++ .../datalake_fdw/datalake_fdw_test.control | 24 + contrib/datalake_fdw/exports.txt | 7 + .../src/am_iceberg/pg_iceberg_extensible.c | 2 + .../src/am_iceberg/pg_iceberg_guc.c | 24 + .../src/am_iceberg/pg_iceberg_guc.h | 1 + contrib/datalake_fdw/src/common/dl_err.c | 12 +- contrib/datalake_fdw/src/common/dl_err.h | 15 +- contrib/datalake_fdw/src/common/dl_resource.c | 139 ++++++ contrib/datalake_fdw/src/common/dl_resource.h | 78 +++ contrib/datalake_fdw/src/common/dl_wrappers.h | 23 +- .../src/common/file_system_wrapper.cpp | 12 +- .../src/common/s3_file_system.cpp | 12 +- .../datalake_fdw/src/format/arrow_builder.cpp | 341 +++++++++++++ .../datalake_fdw/src/format/arrow_builder.h | 95 ++++ .../datalake_fdw/src/format/arrow_decode.c | 355 ++++++++++++++ .../datalake_fdw/src/format/arrow_decode.h | 75 +++ .../datalake_fdw/src/format/arrow_support.cpp | 179 +++++++ .../datalake_fdw/src/format/arrow_support.h | 75 +++ contrib/datalake_fdw/src/format/format.h | 77 ++- .../datalake_fdw/src/format/format_registry.c | 24 +- .../src/format/parquet/parquet_format.cpp | 55 +++ .../src/format/parquet/parquet_format.h | 45 ++ .../src/format/parquet/parquet_internal.h | 47 ++ .../src/format/parquet/parquet_read.cpp | 280 +++++++++++ .../src/format/parquet/parquet_write.cpp | 451 ++++++++++++++++++ .../datalake_fdw/src/test/datalake_fdw_test.c | 411 ++++++++++++++++ .../datalake_fdw/test/automation/README.md | 13 +- .../scripts/test/run_smoke_tests.sh | 14 +- .../expected/parquet_roundtrip.out | 250 ++++++++++ .../format_parquet/sql/parquet_roundtrip.sql | 174 +++++++ 33 files changed, 3429 insertions(+), 42 deletions(-) create mode 100644 contrib/datalake_fdw/datalake_fdw_test--1.0.sql create mode 100644 contrib/datalake_fdw/datalake_fdw_test.control create mode 100644 contrib/datalake_fdw/src/common/dl_resource.c create mode 100644 contrib/datalake_fdw/src/common/dl_resource.h create mode 100644 contrib/datalake_fdw/src/format/arrow_builder.cpp create mode 100644 contrib/datalake_fdw/src/format/arrow_builder.h create mode 100644 contrib/datalake_fdw/src/format/arrow_decode.c create mode 100644 contrib/datalake_fdw/src/format/arrow_decode.h create mode 100644 contrib/datalake_fdw/src/format/arrow_support.cpp create mode 100644 contrib/datalake_fdw/src/format/arrow_support.h create mode 100644 contrib/datalake_fdw/src/format/parquet/parquet_format.cpp create mode 100644 contrib/datalake_fdw/src/format/parquet/parquet_format.h create mode 100644 contrib/datalake_fdw/src/format/parquet/parquet_internal.h create mode 100644 contrib/datalake_fdw/src/format/parquet/parquet_read.cpp create mode 100644 contrib/datalake_fdw/src/format/parquet/parquet_write.cpp create mode 100644 contrib/datalake_fdw/src/test/datalake_fdw_test.c create mode 100644 contrib/datalake_fdw/test/automation/sqlrepo/smoke/format_parquet/expected/parquet_roundtrip.out create mode 100644 contrib/datalake_fdw/test/automation/sqlrepo/smoke/format_parquet/sql/parquet_roundtrip.sql diff --git a/.github/workflows/build-cloudberry.yml b/.github/workflows/build-cloudberry.yml index b790d3411ec..ef84e33dc7f 100644 --- a/.github/workflows/build-cloudberry.yml +++ b/.github/workflows/build-cloudberry.yml @@ -1463,6 +1463,34 @@ jobs: exit 0 fi + # datalake_fdw needs the Arrow and Parquet C++ libraries, which the + # build image does not carry -- nothing in the RPM uses them, so they + # would be weight every other job paid for. + if [[ "${PGXS_EXTENSION}" == "contrib/datalake_fdw" ]]; then + . /etc/os-release + if [[ "${VERSION_ID%%.*}" == "8" ]]; then + # EPEL 8 has them, but its libarrow-devel needs a utf8proc-devel + # that modular filtering keeps out of PowerTools, so it cannot be + # installed. The Arrow project's own repository can. Pinned to + # the version EPEL 10 carries, both because that is one version + # fewer to have working and because the newest wants C++20, which + # Rocky 8's gcc 8 does not have. + # EPEL as well, and not only for Arrow itself: arrow-devel needs + # re2-devel and parquet-devel needs thrift-devel, and on EL8 both + # of those live in EPEL. + dnf install -y \ + https://apache.jfrog.io/artifactory/arrow/almalinux/8/apache-arrow-release-latest.rpm + dnf install -y --enablerepo=epel --enablerepo=powertools \ + arrow-devel-17.0.0-1.el8 parquet-devel-17.0.0-1.el8 + else + # From EPEL, which the image has enrolled but left disabled, + # exactly as it does for its own EPEL packages; CRB carries what + # they depend on. + dnf install -y --enablerepo=epel --enablerepo=crb \ + libarrow-devel parquet-libs-devel + fi + fi + # The RPM installs as root; the build runs as gpadmin, as everywhere # else in this job. -H follows the command-line symlink. chown -RH gpadmin:gpadmin "${BUILD_DESTINATION}/" diff --git a/contrib/datalake_fdw/Makefile b/contrib/datalake_fdw/Makefile index bd1d0179513..0653764581c 100644 --- a/contrib/datalake_fdw/Makefile +++ b/contrib/datalake_fdw/Makefile @@ -18,8 +18,11 @@ # contrib/datalake_fdw/Makefile MODULE_big = datalake_fdw -EXTENSION = datalake_fdw -DATA = datalake_fdw--1.0.sql + +# A second extension, so that installing datalake_fdw does not put the test +# entry points in a production database. One library, so they reach internals. +EXTENSION = datalake_fdw datalake_fdw_test +DATA = datalake_fdw--1.0.sql datalake_fdw_test--1.0.sql OBJS = \ src/am_iceberg/pg_iceberg_am_handler.o \ @@ -36,12 +39,46 @@ OBJS = \ src/meta/meta_engine_init.o \ src/meta/engine_stub/stub_engine.o \ src/format/format_registry.o \ + src/format/arrow_support.o \ + src/format/arrow_builder.o \ + src/format/arrow_decode.o \ + src/format/parquet/parquet_format.o \ + src/format/parquet/parquet_read.o \ + src/format/parquet/parquet_write.o \ src/common/dl_err.o \ + src/common/dl_resource.o \ src/common/dl_option_util.o \ src/common/parser_option.o \ src/common/file_system_wrapper.o \ src/common/s3_file_system.o \ - src/common/backend_registry.o + src/common/backend_registry.o \ + src/test/datalake_fdw_test.o + +# libparquet is written in terms of Arrow's types, so linking one links both. +PKG_CONFIG ?= pkg-config +ARROW_MODULES = arrow parquet +HAVE_ARROW := $(shell $(PKG_CONFIG) --exists $(ARROW_MODULES) 2>/dev/null && echo yes) +# Without the filter, whichever C++ standard Arrow's .pc file names wins: these +# land in CPPFLAGS, which pgxs.mk puts after CXXFLAGS on the command line. The +# Arrow project's own packages say -std=c++11, and their headers then fail to +# compile against themselves. +ARROW_CPPFLAGS := $(filter-out -std=%,\ + $(shell $(PKG_CONFIG) --cflags $(ARROW_MODULES) 2>/dev/null)) +ARROW_LIBS := $(shell $(PKG_CONFIG) --libs $(ARROW_MODULES) 2>/dev/null) + +# This has to refuse at parse time. A recipe hung off `all` would run after +# pgxs.mk's own all-lib, so the compiler would fail on a missing arrow/api.h +# first and this would never be reached. Not for the clean targets, because +# contrib/Makefile recurses here for those even when the module is not +# configured in. +ifneq ($(HAVE_ARROW),yes) +ifeq ($(filter clean distclean maintainer-clean,$(MAKECMDGOALS)),) +$(error datalake_fdw needs the Apache Arrow and Parquet C++ libraries, and \ +pkg-config found neither "arrow" nor "parquet". They are libarrow-devel and \ +parquet-libs-devel on Rocky and RHEL, libarrow-dev and libparquet-dev on \ +Debian and Ubuntu) +endif +endif # Use the documented PGXS knobs: pgxs.mk appends these AFTER the flags configure # chose, so optimization/warning settings survive. A pre-include @@ -49,7 +86,7 @@ OBJS = \ # Makefile.global's own "CFLAGS = @CFLAGS@" assignment. PG_CFLAGS = -fvisibility=hidden PG_CXXFLAGS = -fvisibility=hidden -fvisibility-inlines-hidden -std=c++17 -PG_CPPFLAGS = -I$(srcdir)/src +PG_CPPFLAGS = -I$(srcdir)/src $(ARROW_CPPFLAGS) # The regression cases live with the rest of the test material rather than in a # second place of their own; pg_regress is pointed at them. REGRESS_OPTS is @@ -65,6 +102,10 @@ REGRESS = iceberg_am_ddl iceberg_am_reject iceberg_am_acl REGRESS_OPTS = --temp-config=$(srcdir)/datalake_fdw.conf \ --inputdir=$(srcdir)/test/automation/sqlrepo/smoke/iceberg_am +# A second category, and pg_regress takes one --inputdir, so it is a second run. +FORMAT_PARQUET_REGRESS = parquet_roundtrip +FORMAT_PARQUET_INPUTDIR = $(srcdir)/test/automation/sqlrepo/smoke/format_parquet + EXTRA_CLEAN = exports_darwin.list exports.map # Keep the aggregate target as make's default goal. @@ -87,10 +128,10 @@ endif # Shared libraries are linked with $(CC) (see src/Makefile.shlib COMPILER), so a # module containing C++ translation units must pull in the C++ runtime itself. -SHLIB_LINK += -lstdc++ +SHLIB_LINK += -lstdc++ $(ARROW_LIBS) -# Arrow and other C++ dependencies land in this module later; the export list is -# the single place that decides what stays visible, so the mechanism ships now. +# The export list is the single place that decides what stays visible -- which +# now also means none of Arrow's symbols become symbols this module offers. ifeq ($(PORTNAME), darwin) EXPORT_LIST = exports_darwin.list SHLIB_LINK += -Wl,-exported_symbols_list,exports_darwin.list @@ -107,3 +148,30 @@ endif all: $(EXPORT_LIST) $(shlib): $(EXPORT_LIST) + +# Hung off check and installcheck so that both get both categories. REGRESS_OPTS +# has to come along: pgxs.mk is where --dbname=$(CONTRIB_TESTDB) is added to it, +# and without that pg_regress falls back to "regression" -- which it DROPs and +# recreates, taking the core suite's database with it. The second --inputdir +# wins over the one in REGRESS_OPTS. submake and REGRESS_PREP are the same +# prerequisites pgxs.mk gives its own targets, so that a parallel make cannot +# start pg_regress before it has been built. +installcheck: installcheck-format-parquet + +installcheck-format-parquet: submake $(REGRESS_PREP) + $(pg_regress_installcheck) $(REGRESS_OPTS) \ + --inputdir=$(FORMAT_PARQUET_INPUTDIR) $(FORMAT_PARQUET_REGRESS) + +.PHONY: installcheck-format-parquet + +# "make check" is in-tree only -- under PGXS pgxs.mk refuses the target -- and +# it is the only run that supplies the temp-config that preloads this module. +ifndef USE_PGXS +check: check-format-parquet + +check-format-parquet: submake $(REGRESS_PREP) + $(pg_regress_check) $(REGRESS_OPTS) \ + --inputdir=$(FORMAT_PARQUET_INPUTDIR) $(FORMAT_PARQUET_REGRESS) + +.PHONY: check-format-parquet +endif diff --git a/contrib/datalake_fdw/datalake_fdw_test--1.0.sql b/contrib/datalake_fdw/datalake_fdw_test--1.0.sql new file mode 100644 index 00000000000..bb76a5023ba --- /dev/null +++ b/contrib/datalake_fdw/datalake_fdw_test--1.0.sql @@ -0,0 +1,51 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + * + * contrib/datalake_fdw/datalake_fdw_test--1.0.sql + */ + +-- complain if script is sourced in psql, rather than via CREATE EXTENSION +\echo Use "CREATE EXTENSION datalake_fdw_test" to load this file. \quit + +/* + * Both functions name a path on the server's file system and run as the + * operating system user the server does, so they are as privileged as + * pg_read_server_files and are granted the same way: to nobody, until someone + * decides otherwise. + * + * The reader is pinned to the coordinator. Without that the planner may put a + * function scan on the segments, where each of them would read the whole file + * and the rows would come back as many times as there are segments. The writer + * cannot say the same -- EXECUTE ON is only accepted for a set-returning + * function -- but it does not need to: it is called in a target list with no + * FROM clause, which is evaluated on the coordinator, and the query it runs is + * dispatched from there like any other. + */ +CREATE FUNCTION datalake_parquet_write(path text, + query text, + row_group_size int DEFAULT 0) +RETURNS bigint AS 'MODULE_PATHNAME' LANGUAGE C STRICT VOLATILE; + +REVOKE EXECUTE ON FUNCTION datalake_parquet_write(text, text, int) FROM PUBLIC; + +CREATE FUNCTION datalake_parquet_read(path text, + first_row_group int DEFAULT 0, + n_row_groups int DEFAULT 0) +RETURNS SETOF record AS 'MODULE_PATHNAME' LANGUAGE C STRICT EXECUTE ON COORDINATOR; + +REVOKE EXECUTE ON FUNCTION datalake_parquet_read(text, int, int) FROM PUBLIC; diff --git a/contrib/datalake_fdw/datalake_fdw_test.control b/contrib/datalake_fdw/datalake_fdw_test.control new file mode 100644 index 00000000000..8df236864f0 --- /dev/null +++ b/contrib/datalake_fdw/datalake_fdw_test.control @@ -0,0 +1,24 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# contrib/datalake_fdw/datalake_fdw_test.control + +comment = 'entry points into datalake_fdw internals, for testing' +default_version = '1.0' +module_pathname = '$libdir/datalake_fdw' +relocatable = false +requires = 'datalake_fdw' diff --git a/contrib/datalake_fdw/exports.txt b/contrib/datalake_fdw/exports.txt index 0db251366df..e89c850b512 100644 --- a/contrib/datalake_fdw/exports.txt +++ b/contrib/datalake_fdw/exports.txt @@ -30,3 +30,10 @@ pg_finfo_iceberg_catalog_fdw_validator iceberg_catalog_fdw_validator pg_finfo_iceberg_volume_fdw_validator iceberg_volume_fdw_validator + +# datalake_fdw_test: not part of what this module offers, but a SQL-callable +# function has to be found by name in the library like any other. +pg_finfo_datalake_parquet_write +datalake_parquet_write +pg_finfo_datalake_parquet_read +datalake_parquet_read diff --git a/contrib/datalake_fdw/src/am_iceberg/pg_iceberg_extensible.c b/contrib/datalake_fdw/src/am_iceberg/pg_iceberg_extensible.c index f675a1ef24d..26e1f7a68c5 100644 --- a/contrib/datalake_fdw/src/am_iceberg/pg_iceberg_extensible.c +++ b/contrib/datalake_fdw/src/am_iceberg/pg_iceberg_extensible.c @@ -33,6 +33,7 @@ #include "access/table.h" #include "access/tableam.h" #include "am_iceberg/pg_iceberg_ddl.h" +#include "common/dl_resource.h" #include "am_iceberg/pg_iceberg_guc.h" #include "am_iceberg/pg_iceberg_options.h" #include "am_iceberg/pg_iceberg_reject.h" @@ -951,6 +952,7 @@ _PG_init(void) errmsg("datalake_fdw must be loaded via shared_preload_libraries"), errhint("Add \"datalake_fdw\" to shared_preload_libraries and restart the server."))); + dl_resource_init(); pg_iceberg_define_gucs(); pg_iceberg_register_reloptions(); DatalakeRegisterMetaEngines(); diff --git a/contrib/datalake_fdw/src/am_iceberg/pg_iceberg_guc.c b/contrib/datalake_fdw/src/am_iceberg/pg_iceberg_guc.c index 71dc5c53ef1..5e9ec353f64 100644 --- a/contrib/datalake_fdw/src/am_iceberg/pg_iceberg_guc.c +++ b/contrib/datalake_fdw/src/am_iceberg/pg_iceberg_guc.c @@ -33,6 +33,7 @@ char *iceberg_default_catalog; char *iceberg_default_volume; +int iceberg_batch_rows; void pg_iceberg_define_gucs(void) @@ -64,4 +65,27 @@ pg_iceberg_define_gucs(void) NULL, NULL, NULL); + + /* + * How many rows travel between the executor and a data file at a time. + * Every per-batch cost is paid once per this many rows, and the batch and + * its Arrow copy are held while it is built, so the right value trades + * memory for that -- which depends on how wide the table is, and is why + * this is a setting rather than a constant. + * + * The ceiling is Parquet's default row group length: a batch bigger than + * the unit a file is written in buys nothing. + */ + DefineCustomIntVariable("iceberg.batch_rows", + "Rows per batch exchanged with a lake table's data files.", + NULL, + &iceberg_batch_rows, + 16384, + 1, + 1024 * 1024, + PGC_USERSET, + 0, + NULL, + NULL, + NULL); } diff --git a/contrib/datalake_fdw/src/am_iceberg/pg_iceberg_guc.h b/contrib/datalake_fdw/src/am_iceberg/pg_iceberg_guc.h index 8cd8d0a4400..68d30bb1908 100644 --- a/contrib/datalake_fdw/src/am_iceberg/pg_iceberg_guc.h +++ b/contrib/datalake_fdw/src/am_iceberg/pg_iceberg_guc.h @@ -31,6 +31,7 @@ extern char *iceberg_default_catalog; extern char *iceberg_default_volume; +extern int iceberg_batch_rows; extern void pg_iceberg_define_gucs(void); diff --git a/contrib/datalake_fdw/src/common/dl_err.c b/contrib/datalake_fdw/src/common/dl_err.c index 4780d9d44a0..44f32cf307b 100644 --- a/contrib/datalake_fdw/src/common/dl_err.c +++ b/contrib/datalake_fdw/src/common/dl_err.c @@ -160,10 +160,20 @@ dl_err_message(DlErrCode code) void dl_error_report(int elevel, DlErrCode code, const char *prefix) { - const DlErrorDetail *detail = dl_error_get(); + /* + * A copy, because reporting consumes the record: what is left otherwise is + * a description of a failure that has already been reported, waiting for + * the next failure with the same code to adopt it -- and the check below + * cannot tell those two apart. The reset has to happen before the ereport, + * which at ERROR does not come back. + */ + DlErrorDetail detail_copy = *dl_error_get(); + const DlErrorDetail *detail = &detail_copy; StringInfoData detail_buf; bool has_detail; + dl_error_reset(); + /* * Detail recorded against a different code belongs to some other failure -- * an implementation that reported this one without recording anything, for diff --git a/contrib/datalake_fdw/src/common/dl_err.h b/contrib/datalake_fdw/src/common/dl_err.h index 67949644c95..7d0e652ed4d 100644 --- a/contrib/datalake_fdw/src/common/dl_err.h +++ b/contrib/datalake_fdw/src/common/dl_err.h @@ -109,10 +109,23 @@ extern const char *dl_err_message(DlErrCode code); * when the session asked for log-level detail -- a stack is for whoever is * debugging the implementation, not for whoever ran the statement. * - * Detail recorded against a different code is ignored rather than misattributed. + * Detail recorded against a different code is ignored rather than misattributed, + * and reporting consumes what it used. Matching on the code alone cannot tell + * this failure's detail from an earlier failure's with the same code, so the + * record is not left behind for the next one to inherit. */ extern void dl_error_report(int elevel, DlErrCode code, const char *prefix); +/* + * An argument that was not what it had to be. A caller reporting the code + * alone would say "invalid parameter" about a call the user never wrote, so + * this names the entry point instead -- there is no user error to describe, + * only which one of ours was called wrongly. + */ +#define DL_ARG_ERROR(operation) \ + (dl_error_set(DL_ERR_INVALID_OPTION, (operation), NULL, \ + "a required argument was missing"), DL_ERR_INVALID_OPTION) + #ifdef __cplusplus } #endif diff --git a/contrib/datalake_fdw/src/common/dl_resource.c b/contrib/datalake_fdw/src/common/dl_resource.c new file mode 100644 index 00000000000..691180c2407 --- /dev/null +++ b/contrib/datalake_fdw/src/common/dl_resource.c @@ -0,0 +1,139 @@ +/*------------------------------------------------------------------------- + * + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + * + * dl_resource.c + * Cleanups that happen even when nothing calls them. + * + * IDENTIFICATION + * contrib/datalake_fdw/src/common/dl_resource.c + * + *------------------------------------------------------------------------- + */ + +#include "postgres.h" + +#include + +#include "storage/ipc.h" +#include "utils/resowner.h" + +#include "common/dl_resource.h" + +typedef struct DlResourceEntry +{ + struct DlResourceEntry *next; + ResourceOwner owner; + DlResourceRelease release; + void *arg; +} DlResourceEntry; + +/* + * There are a handful of these at a time -- one per open data file -- so a list + * walked linearly is the whole structure needed. + */ +static DlResourceEntry *dl_resources; + +/* + * malloc rather than palloc: this outlives the memory context that was current + * when it was remembered, by construction, and is walked while the transaction + * is being torn down. + */ +static void +dl_resource_release_callback(ResourceReleasePhase phase, bool isCommit, + bool isTopLevel, void *arg) +{ + DlResourceEntry **link; + + /* + * After locks, so that anything the release path might touch is still + * usable. Nothing to do while the process is exiting: the descriptors go + * with it, and running C++ destructors on the way out is a way to turn an + * exit into a crash. + */ + if (phase != RESOURCE_RELEASE_AFTER_LOCKS || proc_exit_inprogress) + return; + + link = &dl_resources; + while (*link != NULL) + { + DlResourceEntry *entry = *link; + + if (entry->owner != CurrentResourceOwner) + { + link = &entry->next; + continue; + } + + /* + * Reaching here on a commit means the owner released nothing: the + * statement finished and left a file open. The resource is still + * cleaned up, but quietly doing so would hide the bug that let it + * happen, which is the same call PostgreSQL's own resource owners make. + */ + if (isCommit) + elog(WARNING, "datalake_fdw leaked a resource: %p", entry->arg); + + *link = entry->next; + entry->release(entry->arg); + free(entry); + } +} + +void +dl_resource_init(void) +{ + RegisterResourceReleaseCallback(dl_resource_release_callback, NULL); +} + +bool +dl_resource_remember(DlResourceRelease release, void *arg) +{ + DlResourceEntry *entry = malloc(sizeof(DlResourceEntry)); + + if (entry == NULL) + return false; + + entry->owner = CurrentResourceOwner; + entry->release = release; + entry->arg = arg; + entry->next = dl_resources; + dl_resources = entry; + + return true; +} + +void +dl_resource_forget(DlResourceRelease release, void *arg) +{ + DlResourceEntry **link = &dl_resources; + + while (*link != NULL) + { + DlResourceEntry *entry = *link; + + if (entry->release == release && entry->arg == arg) + { + *link = entry->next; + free(entry); + return; + } + + link = &entry->next; + } +} diff --git a/contrib/datalake_fdw/src/common/dl_resource.h b/contrib/datalake_fdw/src/common/dl_resource.h new file mode 100644 index 00000000000..00ecc51c7d8 --- /dev/null +++ b/contrib/datalake_fdw/src/common/dl_resource.h @@ -0,0 +1,78 @@ +/*------------------------------------------------------------------------- + * + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + * + * dl_resource.h + * Cleanups that happen even when nothing calls them. + * + * A reader or a writer holds an operating system file descriptor and memory + * from an allocator that is not PostgreSQL's. Neither is reclaimed by + * transaction abort, so anything that owns one has to be released by name -- + * and a caller that raises an error before it reaches its own cleanup would + * never get to. Registering here makes the resource owner do it instead, and + * makes a caller that simply forgot a warning rather than a descriptor that is + * gone until the backend exits. + * + * This is the shape PAX uses (contrib/pax_storage/src/cpp/comm/pax_resource.cc): + * a callback registered once, and a list of what to release keyed by the owner + * that was current when it was remembered. PostgreSQL 16 also has a typed + * resource-kind API, which this server does not carry. + * + * IDENTIFICATION + * contrib/datalake_fdw/src/common/dl_resource.h + * + *------------------------------------------------------------------------- + */ + +#ifndef DL_RESOURCE_H +#define DL_RESOURCE_H + +#ifdef __cplusplus +extern "C" +{ +#endif + +/* + * Runs during transaction abort, so the same rules as any other cleanup path: + * noexcept, and it must not raise. + */ +typedef void (*DlResourceRelease) (void *arg); + +/* Installs the callback. Called once, from _PG_init. */ +extern void dl_resource_init(void); + +/* + * Remembers that `release(arg)` has to happen before the current resource owner + * goes away. Returns false only when it could not record that, which the + * caller has to treat as a failure to acquire the resource at all -- releasing + * it itself and reporting -- because nothing else is going to. + */ +extern bool dl_resource_remember(DlResourceRelease release, void *arg); + +/* + * Drops that record, for the ordinary path where the caller releases the + * resource itself. Silent when there is nothing to drop: a cleanup that runs + * twice reaches this the second time. + */ +extern void dl_resource_forget(DlResourceRelease release, void *arg); + +#ifdef __cplusplus +} +#endif + +#endif /* DL_RESOURCE_H */ diff --git a/contrib/datalake_fdw/src/common/dl_wrappers.h b/contrib/datalake_fdw/src/common/dl_wrappers.h index 737d97fae7d..5aadc9b4bce 100644 --- a/contrib/datalake_fdw/src/common/dl_wrappers.h +++ b/contrib/datalake_fdw/src/common/dl_wrappers.h @@ -136,15 +136,36 @@ dl_can_log_cleanup_warning(void) errmsg("datalake_fdw: %s", dl_error_msg_))); \ } while (0) +/* + * Class 2 guards record what happened as well as that it happened. + * + * The detail is one per-backend record, and dl_error_report() decides whether + * it belongs to the failure being reported by comparing codes -- which cannot + * tell this DL_ERR_INTERNAL from an earlier one. A guard that set the code and + * recorded nothing would therefore report the previous statement's message as + * the cause of this one. So it always records, and what a C++ exception has to + * say is the only description of it there is going to be. PAX takes the same + * line in CBDB_END_TRY(), where an unnamed failure falls back to the function + * it happened in rather than to whatever was there before. + * + * `operation` names the call, the way the metadata engine's dispatch does. + */ #define DL_ABI_GUARD_BEGIN \ try \ { -#define DL_ABI_GUARD_END(errvar) \ +#define DL_ABI_GUARD_END(errvar, operation) \ + } \ + catch (const std::exception &e) \ + { \ + (errvar) = DL_ERR_INTERNAL; \ + dl_error_set(DL_ERR_INTERNAL, (operation), NULL, e.what()); \ } \ catch (...) \ { \ (errvar) = DL_ERR_INTERNAL; \ + dl_error_set(DL_ERR_INTERNAL, (operation), NULL, \ + "unknown C++ exception"); \ } #define DL_CLEANUP_GUARD_BEGIN \ diff --git a/contrib/datalake_fdw/src/common/file_system_wrapper.cpp b/contrib/datalake_fdw/src/common/file_system_wrapper.cpp index 2b576fe4344..c0c8b2d3e51 100644 --- a/contrib/datalake_fdw/src/common/file_system_wrapper.cpp +++ b/contrib/datalake_fdw/src/common/file_system_wrapper.cpp @@ -71,7 +71,7 @@ datalake_fs_open(const DatalakeLocation *location, } } } - DL_ABI_GUARD_END(rc); + DL_ABI_GUARD_END(rc, "fs_open"); return rc; } @@ -116,7 +116,7 @@ datalake_fs_list(DatalakeFileSystem fs, const char *prefix, rc = fs->ops->fs_list(fs, prefix, names_out, nnames_out); } } - DL_ABI_GUARD_END(rc); + DL_ABI_GUARD_END(rc, "fs_list"); return rc; } @@ -144,7 +144,7 @@ datalake_file_open(DatalakeFileSystem fs, const char *path, (*file_out)->ops = fs->ops; } } - DL_ABI_GUARD_END(rc); + DL_ABI_GUARD_END(rc, "file_open"); return rc; } @@ -165,7 +165,7 @@ datalake_file_read(DatalakeFile file, void *buffer, int64_t length, rc = file->ops->file_read(file, buffer, length, nread); } } - DL_ABI_GUARD_END(rc); + DL_ABI_GUARD_END(rc, "file_read"); return rc; } @@ -182,7 +182,7 @@ datalake_file_write(DatalakeFile file, const void *buffer, int64_t length) else rc = file->ops->file_write(file, buffer, length); } - DL_ABI_GUARD_END(rc); + DL_ABI_GUARD_END(rc, "file_write"); return rc; } @@ -209,7 +209,7 @@ datalake_file_close(DatalakeFile *file) rc = doomed->ops->file_close(doomed); } } - DL_ABI_GUARD_END(rc); + DL_ABI_GUARD_END(rc, "file_close"); return rc; } diff --git a/contrib/datalake_fdw/src/common/s3_file_system.cpp b/contrib/datalake_fdw/src/common/s3_file_system.cpp index d255200528b..ebb7fadf2ad 100644 --- a/contrib/datalake_fdw/src/common/s3_file_system.cpp +++ b/contrib/datalake_fdw/src/common/s3_file_system.cpp @@ -122,7 +122,7 @@ s3_fs_open(const DatalakeLocation *location, const DlKeyValue *credentials, *fs_out = handle.release(); } } - DL_ABI_GUARD_END(rc); + DL_ABI_GUARD_END(rc, "s3_fs_open"); return rc; } @@ -153,7 +153,7 @@ s3_fs_list(DatalakeFileSystem fs, const char *prefix, char ***names_out, else rc = handle->impl->List(prefix, names_out, nnames_out); } - DL_ABI_GUARD_END(rc); + DL_ABI_GUARD_END(rc, "s3_fs_list"); return rc; } @@ -173,7 +173,7 @@ s3_file_open(DatalakeFileSystem fs, const char *path, DatalakeFileMode mode, else rc = handle->impl->OpenFile(path, mode, file_out); } - DL_ABI_GUARD_END(rc); + DL_ABI_GUARD_END(rc, "s3_file_open"); return rc; } @@ -194,7 +194,7 @@ s3_file_read(DatalakeFile file, void *buffer, int64_t length, int64_t *nread) rc = DL_ERR_NOT_SUPPORTED; } - DL_ABI_GUARD_END(rc); + DL_ABI_GUARD_END(rc, "s3_file_read"); return rc; } @@ -212,7 +212,7 @@ s3_file_write(DatalakeFile file, const void *buffer, int64_t length) rc = DL_ERR_NOT_SUPPORTED; } - DL_ABI_GUARD_END(rc); + DL_ABI_GUARD_END(rc, "s3_file_write"); return rc; } @@ -228,7 +228,7 @@ s3_file_close(DatalakeFile file) rc = DL_ERR_NOT_SUPPORTED; } - DL_ABI_GUARD_END(rc); + DL_ABI_GUARD_END(rc, "s3_file_close"); return rc; } diff --git a/contrib/datalake_fdw/src/format/arrow_builder.cpp b/contrib/datalake_fdw/src/format/arrow_builder.cpp new file mode 100644 index 00000000000..3dd68f4006a --- /dev/null +++ b/contrib/datalake_fdw/src/format/arrow_builder.cpp @@ -0,0 +1,341 @@ +/*------------------------------------------------------------------------- + * + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + * + * arrow_builder.cpp + * Accumulation of PostgreSQL tuples into an Arrow batch. + * + * IDENTIFICATION + * contrib/datalake_fdw/src/format/arrow_builder.cpp + * + *------------------------------------------------------------------------- + */ + +#include +#include + +#include +#include + +#include "format/arrow_support.h" + +#include "common/dl_resource.h" +#include "common/dl_wrappers.h" +#include "format/arrow_builder.h" + +extern "C" +{ +#include "catalog/pg_type.h" +#include "utils/date.h" +#include "utils/timestamp.h" +#include "varatt.h" +} + +/* + * PostgreSQL counts from 2000-01-01 and Arrow from 1970-01-01. Everything + * below that touches a date or a timestamp shifts by this, and getting the sign + * wrong is a 30-year error that no round trip through our own code would + * notice -- both halves would agree. It is written once, here. + */ +#define DL_EPOCH_DELTA_DAYS ((int32) (POSTGRES_EPOCH_JDATE - UNIX_EPOCH_JDATE)) +#define DL_EPOCH_DELTA_USECS (((int64) DL_EPOCH_DELTA_DAYS) * USECS_PER_DAY) + +struct DlArrowBuilderData +{ + std::shared_ptr schema; + std::vector types; + std::vector> builders; + int64_t nrows; +}; + +/* + * One value into the builder for its column. The Datum has already been + * detoasted by the caller, so nothing here allocates and nothing here can + * raise. + */ +static arrow::Status +dl_append_datum(arrow::ArrayBuilder *builder, Oid atttypid, Datum value) +{ + switch (atttypid) + { + case BOOLOID: + return static_cast(builder) + ->Append(DatumGetBool(value)); + case INT2OID: + return static_cast(builder) + ->Append(DatumGetInt16(value)); + case INT4OID: + return static_cast(builder) + ->Append(DatumGetInt32(value)); + case INT8OID: + return static_cast(builder) + ->Append(DatumGetInt64(value)); + case FLOAT4OID: + return static_cast(builder) + ->Append(DatumGetFloat4(value)); + case FLOAT8OID: + return static_cast(builder) + ->Append(DatumGetFloat8(value)); + + case TEXTOID: + case VARCHAROID: + case BPCHAROID: + { + struct varlena *v = (struct varlena *) DatumGetPointer(value); + + return static_cast(builder) + ->Append(VARDATA_ANY(v), VARSIZE_ANY_EXHDR(v)); + } + + case BYTEAOID: + { + struct varlena *v = (struct varlena *) DatumGetPointer(value); + + return static_cast(builder) + ->Append(VARDATA_ANY(v), VARSIZE_ANY_EXHDR(v)); + } + + case DATEOID: + { + DateADT date = DatumGetDateADT(value); + + /* + * Parquet has no way to say "infinity", and writing the + * sentinel would hand the next reader a date 5.8 million years + * out as if it were a real one. + */ + if (DATE_NOT_FINITE(date)) + return arrow::Status::NotImplemented( + "an infinite date cannot be written to a data file"); + + return static_cast(builder) + ->Append(date + DL_EPOCH_DELTA_DAYS); + } + + case TIMESTAMPOID: + case TIMESTAMPTZOID: + { + Timestamp ts = DatumGetTimestamp(value); + + if (TIMESTAMP_NOT_FINITE(ts)) + return arrow::Status::NotImplemented( + "an infinite timestamp cannot be written to a data file"); + + /* + * PostgreSQL's range runs about 34 years past the last instant + * Arrow can hold in microseconds from 1970, so the shift below + * is not always representable. Without this the addition wraps + * -- quietly, because the build sets -fwrapv -- and a year + * 294250 timestamp is written as one 292000 years before the + * epoch, with the write reporting success. The read side has + * the mirror of this guard, and neither can stand in for the + * other: a round trip through both would agree. + */ + if (ts > PG_INT64_MAX - DL_EPOCH_DELTA_USECS) + return arrow::Status::Invalid( + "timestamp is too far in the future to be written to a " + "data file"); + + return static_cast(builder) + ->Append(ts + DL_EPOCH_DELTA_USECS); + } + + default: + + /* + * Unreachable until a new type is added to one of the two + * switches and not the other, which is why it names the OID. A + * number and not a name: resolving one means a catalog lookup, and + * nothing on this side of the ABI may allocate or raise. + */ + return arrow::Status::NotImplemented( + "no Arrow type is mapped for PostgreSQL type OID ", atttypid); + } +} + +/* What the resource owner calls if nothing else did. */ +extern "C" void +dl_arrow_builder_release(void *arg) +{ + DL_CLEANUP_GUARD_BEGIN + { + delete static_cast(arg); + } + DL_CLEANUP_GUARD_END; +} + +extern "C" DlErrCode +dl_arrow_builder_open(void *tupdesc_arg, DlArrowBuilder *out) +{ + DlErrCode result = DL_OK; + + if (out == NULL) + return DL_ARG_ERROR("open_builder"); + *out = NULL; + + if (tupdesc_arg == NULL) + return DL_ARG_ERROR("open_builder"); + + DL_ABI_GUARD_BEGIN + { + TupleDesc tupdesc = (TupleDesc) tupdesc_arg; + std::unique_ptr builder(new DlArrowBuilderData()); + + builder->schema = DlArrowSchemaFromTupleDesc(tupdesc); + if (builder->schema == nullptr) + return DL_ERR_NOT_SUPPORTED; /* detail already recorded */ + + builder->nrows = 0; + builder->types.reserve(tupdesc->natts); + builder->builders.reserve(tupdesc->natts); + + for (int i = 0; i < tupdesc->natts; i++) + { + std::unique_ptr column; + arrow::Status status = arrow::MakeBuilder(arrow::default_memory_pool(), + builder->schema->field(i)->type(), + &column); + + if (!status.ok()) + return DlArrowStatus(status, "create an Arrow array builder"); + + builder->types.push_back(TupleDescAttr(tupdesc, i)->atttypid); + builder->builders.push_back(std::move(column)); + } + + /* + * The buffers behind the builder come from Arrow's allocator, which + * transaction abort knows nothing about. Last thing that may fail. + */ + if (!dl_resource_remember(dl_arrow_builder_release, builder.get())) + { + dl_error_set(DL_ERR_INTERNAL, "open_builder", NULL, + "could not record the batch builder for cleanup"); + return DL_ERR_INTERNAL; + } + + *out = builder.release(); + } + DL_ABI_GUARD_END(result, "open_builder"); + + return result; +} + +extern "C" DlErrCode +dl_arrow_builder_append(DlArrowBuilder builder, const Datum *values, + const bool *nulls, int nvalues) +{ + DlErrCode result = DL_OK; + + if (builder == NULL || values == NULL || nulls == NULL) + return DL_ARG_ERROR("append_row"); + + if (nvalues != (int) builder->builders.size()) + { + dl_error_set(DL_ERR_INTERNAL, "append an Arrow row", NULL, + "the row has a different number of columns than the batch"); + return DL_ERR_INTERNAL; + } + + DL_ABI_GUARD_BEGIN + { + for (int i = 0; i < nvalues; i++) + { + arrow::Status status = nulls[i] + ? builder->builders[i]->AppendNull() + : dl_append_datum(builder->builders[i].get(), builder->types[i], + values[i]); + + if (!status.ok()) + return DlArrowStatus(status, "append a value to an Arrow array"); + } + + builder->nrows++; + } + DL_ABI_GUARD_END(result, "append_row"); + + return result; +} + +extern "C" int64_t +dl_arrow_builder_nrows(DlArrowBuilder builder) +{ + return builder == NULL ? 0 : builder->nrows; +} + +extern "C" DlErrCode +dl_arrow_builder_flush(DlArrowBuilder builder, struct ArrowArray *out) +{ + DlErrCode result = DL_OK; + + if (builder == NULL || out == NULL) + return DL_ARG_ERROR("build_batch"); + + DL_ABI_GUARD_BEGIN + { + std::vector> columns; + + columns.reserve(builder->builders.size()); + + for (auto &column : builder->builders) + { + std::shared_ptr array; + + /* Finish() also resets the builder, so the next batch starts here. */ + arrow::Status status = column->Finish(&array); + + if (!status.ok()) + return DlArrowStatus(status, "finish an Arrow array"); + + columns.push_back(std::move(array)); + } + + std::shared_ptr batch = + arrow::RecordBatch::Make(builder->schema, builder->nrows, columns); + + /* + * The schema travels with the writer, which was opened from the same + * descriptor, so exporting it with every batch would be a copy nobody + * reads. + */ + arrow::Status status = arrow::ExportRecordBatch(*batch, out, nullptr); + + if (!status.ok()) + return DlArrowStatus(status, "export an Arrow batch"); + + builder->nrows = 0; + } + DL_ABI_GUARD_END(result, "build_batch"); + + return result; +} + +extern "C" void +dl_arrow_builder_close(DlArrowBuilder *builder) +{ + if (builder == NULL || *builder == NULL) + return; + + /* Cleared first; see the note in parquet_reader_close(). */ + DlArrowBuilderData *impl = *builder; + + *builder = NULL; + dl_resource_forget(dl_arrow_builder_release, impl); + + dl_arrow_builder_release(impl); +} diff --git a/contrib/datalake_fdw/src/format/arrow_builder.h b/contrib/datalake_fdw/src/format/arrow_builder.h new file mode 100644 index 00000000000..d829199beb2 --- /dev/null +++ b/contrib/datalake_fdw/src/format/arrow_builder.h @@ -0,0 +1,95 @@ +/*------------------------------------------------------------------------- + * + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + * + * arrow_builder.h + * Accumulation of PostgreSQL tuples into an Arrow batch. + * + * This is the write half of the boundary the format layer is built on: rows + * arrive one at a time from an executor, and a data file wants them a column at + * a time. Nothing above this knows Arrow, and nothing here knows which format + * the batch ends up in. + * + * postgres.h must be included before this header; the Datum in the append + * signature is the whole reason a tuple can be handed over without copying it + * first. + * + * IDENTIFICATION + * contrib/datalake_fdw/src/format/arrow_builder.h + * + *------------------------------------------------------------------------- + */ + +#ifndef DL_ARROW_BUILDER_H +#define DL_ARROW_BUILDER_H + +#include "common/dl_err.h" +#include "format/format.h" + +#ifdef __cplusplus +extern "C" +{ +#endif + +typedef struct DlArrowBuilderData *DlArrowBuilder; + +/* + * How many rows to accumulate before flushing is the caller's decision, not the + * builder's -- the caller is the one that also has to decide when to roll to a + * new file. iceberg.batch_rows is the setting they take it from. + */ + +/* + * `tupdesc` is a TupleDesc. It is taken as void * so that this header stays + * usable from the C++ side without dragging PostgreSQL's headers into it in a + * particular order; the type is checked by the only two callers there are. + * + * The descriptor has to outlive the builder, which is no constraint in + * practice: the tuples being appended come from it. + */ +extern DlErrCode dl_arrow_builder_open(void *tupdesc, DlArrowBuilder *out); + +/* + * Appends one row. Varlena values must already be detoasted -- this runs on + * the C++ side, where a PostgreSQL error would unwind through frames that + * cannot handle one, so it does not call anything that allocates. + */ +extern DlErrCode dl_arrow_builder_append(DlArrowBuilder builder, + const Datum *values, + const bool *nulls, + int nvalues); + +/* Rows accumulated since the last flush. */ +extern int64_t dl_arrow_builder_nrows(DlArrowBuilder builder); + +/* + * Hands over what has accumulated and starts a new batch. The caller owns the + * exported array and releases it -- or gives it to a writer, which consumes it. + * Flushing nothing is not an error and produces an empty batch. + */ +extern DlErrCode dl_arrow_builder_flush(DlArrowBuilder builder, + struct ArrowArray *out); + +/* Cleanup entry point: releases the builder and clears the caller's handle. */ +extern void dl_arrow_builder_close(DlArrowBuilder *builder); + +#ifdef __cplusplus +} +#endif + +#endif /* DL_ARROW_BUILDER_H */ diff --git a/contrib/datalake_fdw/src/format/arrow_decode.c b/contrib/datalake_fdw/src/format/arrow_decode.c new file mode 100644 index 00000000000..6574eeb2a9b --- /dev/null +++ b/contrib/datalake_fdw/src/format/arrow_decode.c @@ -0,0 +1,355 @@ +/*------------------------------------------------------------------------- + * + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + * + * arrow_decode.c + * PostgreSQL values out of an Arrow batch. + * + * IDENTIFICATION + * contrib/datalake_fdw/src/format/arrow_decode.c + * + *------------------------------------------------------------------------- + */ + +#include "postgres.h" + +#include + +#include "catalog/pg_type.h" +#include "utils/builtins.h" +#include "utils/date.h" +#include "utils/fmgrprotos.h" +#include "utils/timestamp.h" +#include "varatt.h" + +#include "format/arrow_decode.h" + +/* + * The same shift as in arrow_builder.cpp, in the other direction: PostgreSQL + * counts from 2000-01-01 and Arrow from 1970-01-01. + */ +#define DL_EPOCH_DELTA_DAYS ((int32) (POSTGRES_EPOCH_JDATE - UNIX_EPOCH_JDATE)) +#define DL_EPOCH_DELTA_USECS (((int64) DL_EPOCH_DELTA_DAYS) * USECS_PER_DAY) + +/* + * Arrow spells its types as a short string. Only the ones a column of ours can + * be stored as are listed; anything else is a file we did not write, or one + * written by a version that knows more types than this one. + */ +#define DL_ARROW_FORMAT_BOOL "b" +#define DL_ARROW_FORMAT_INT16 "s" +#define DL_ARROW_FORMAT_INT32 "i" +#define DL_ARROW_FORMAT_INT64 "l" +#define DL_ARROW_FORMAT_FLOAT32 "f" +#define DL_ARROW_FORMAT_FLOAT64 "g" +#define DL_ARROW_FORMAT_UTF8 "u" +#define DL_ARROW_FORMAT_BINARY "z" +#define DL_ARROW_FORMAT_DATE32 "tdD" + +/* A timestamp is "tsu:" followed by the time zone, which may be empty. */ +#define DL_ARROW_FORMAT_TIMESTAMP_US "tsu:" + +static DlErrCode +dl_arrow_decode_refuse(const char *arrow_format, Oid atttypid) +{ + char message[256]; + + snprintf(message, sizeof(message), + "a column stored as Arrow type \"%s\" cannot be read as %s", + arrow_format == NULL ? "" : arrow_format, + format_type_be(atttypid)); + + dl_error_set(DL_ERR_NOT_SUPPORTED, "decode an Arrow column", NULL, message); + return DL_ERR_NOT_SUPPORTED; +} + +DlErrCode +dl_arrow_decode_check(const struct ArrowSchema *field, Oid atttypid) +{ + const char *format; + const char *expected; + + if (field == NULL || field->format == NULL) + { + dl_error_set(DL_ERR_INTERNAL, "decode an Arrow column", NULL, + "the batch has a column with no type"); + return DL_ERR_INTERNAL; + } + + format = field->format; + + switch (atttypid) + { + case BOOLOID: + expected = DL_ARROW_FORMAT_BOOL; + break; + case INT2OID: + expected = DL_ARROW_FORMAT_INT16; + break; + case INT4OID: + expected = DL_ARROW_FORMAT_INT32; + break; + case INT8OID: + expected = DL_ARROW_FORMAT_INT64; + break; + case FLOAT4OID: + expected = DL_ARROW_FORMAT_FLOAT32; + break; + case FLOAT8OID: + expected = DL_ARROW_FORMAT_FLOAT64; + break; + case TEXTOID: + case VARCHAROID: + case BPCHAROID: + expected = DL_ARROW_FORMAT_UTF8; + break; + case BYTEAOID: + expected = DL_ARROW_FORMAT_BINARY; + break; + case DATEOID: + expected = DL_ARROW_FORMAT_DATE32; + break; + + case TIMESTAMPOID: + case TIMESTAMPTZOID: + { + const char *zone; + size_t prefix_len = strlen(DL_ARROW_FORMAT_TIMESTAMP_US); + + if (strncmp(format, DL_ARROW_FORMAT_TIMESTAMP_US, prefix_len) != 0) + return dl_arrow_decode_refuse(format, atttypid); + + /* + * Arrow stores a zoned timestamp as the instant in UTC and + * keeps the zone only to display it, so which zone the file + * names does not change the value -- but whether it names one + * at all is the difference between the two PostgreSQL types, + * and reading one as the other would shift every value by the + * session's offset from UTC. + */ + zone = format + prefix_len; + if ((zone[0] != '\0') != (atttypid == TIMESTAMPTZOID)) + return dl_arrow_decode_refuse(format, atttypid); + + return DL_OK; + } + + default: + return dl_arrow_decode_refuse(format, atttypid); + } + + if (strcmp(format, expected) != 0) + return dl_arrow_decode_refuse(format, atttypid); + + return DL_OK; +} + +/* + * Arrow keeps the validity bitmap in the first buffer, and a column with no + * nulls may leave it out entirely. Bit set means present. + */ +static bool +dl_arrow_value_is_null(const struct ArrowArray *column, int64_t row) +{ + const uint8 *validity; + int64 index; + + if (column->n_buffers < 1) + return false; + + validity = (const uint8 *) column->buffers[0]; + if (validity == NULL) + return false; + + index = column->offset + row; + return (validity[index >> 3] & (1 << (index & 7))) == 0; +} + +/* The values buffer of a fixed-width column, already advanced past the offset. */ +#define DL_ARROW_VALUES(column, type) \ + (((const type *) (column)->buffers[1]) + (column)->offset) + +static DlErrCode +dl_arrow_out_of_range(const char *what) +{ + char message[128]; + + snprintf(message, sizeof(message), + "the file holds a %s outside the range PostgreSQL can represent", + what); + + dl_error_set(DL_ERR_INVALID_OPTION, "decode an Arrow column", NULL, message); + return DL_ERR_INVALID_OPTION; +} + +/* + * A variable-length value: an offsets buffer of int32 and one run of bytes. + * Both text and bytea are laid out this way, and differ only in the header the + * copy gets. + */ +static void +dl_arrow_varlen(const struct ArrowArray *column, int64_t row, + const char **data, int32 *length) +{ + const int32 *offsets = DL_ARROW_VALUES(column, int32); + const char *bytes = (const char *) column->buffers[2]; + + *data = bytes + offsets[row]; + *length = offsets[row + 1] - offsets[row]; +} + +DlErrCode +dl_arrow_decode_value(const struct ArrowArray *column, int64_t row, + Oid atttypid, int32 atttypmod, + Datum *value, bool *isnull) +{ + *value = (Datum) 0; + *isnull = true; + + if (row < 0 || row >= column->length) + { + dl_error_set(DL_ERR_INTERNAL, "decode an Arrow column", NULL, + "a row was asked for past the end of the batch"); + return DL_ERR_INTERNAL; + } + + if (dl_arrow_value_is_null(column, row)) + return DL_OK; + + *isnull = false; + + switch (atttypid) + { + case BOOLOID: + { + /* Booleans are a bitmap of their own, not a byte per value. */ + const uint8 *bits = (const uint8 *) column->buffers[1]; + int64 index = column->offset + row; + + *value = BoolGetDatum((bits[index >> 3] & (1 << (index & 7))) != 0); + return DL_OK; + } + + case INT2OID: + *value = Int16GetDatum(DL_ARROW_VALUES(column, int16)[row]); + return DL_OK; + case INT4OID: + *value = Int32GetDatum(DL_ARROW_VALUES(column, int32)[row]); + return DL_OK; + case INT8OID: + *value = Int64GetDatum(DL_ARROW_VALUES(column, int64)[row]); + return DL_OK; + case FLOAT4OID: + *value = Float4GetDatum(DL_ARROW_VALUES(column, float)[row]); + return DL_OK; + case FLOAT8OID: + *value = Float8GetDatum(DL_ARROW_VALUES(column, double)[row]); + return DL_OK; + + case TEXTOID: + case VARCHAROID: + case BPCHAROID: + { + const char *data; + int32 length; + + dl_arrow_varlen(column, row, &data, &length); + *value = PointerGetDatum(cstring_to_text_with_len(data, length)); + + /* + * The file records the bytes and nothing about the length the + * column was declared with, so the value is put through the + * same coercion an inserted one would be: char(n) comes back + * padded to n, and a value too long for a varchar(n) is an + * error rather than something the executor has to meet later. + */ + if (atttypmod >= 0 && atttypid == BPCHAROID) + *value = DirectFunctionCall3(bpchar, *value, + Int32GetDatum(atttypmod), + BoolGetDatum(false)); + else if (atttypmod >= 0 && atttypid == VARCHAROID) + *value = DirectFunctionCall3(varchar, *value, + Int32GetDatum(atttypmod), + BoolGetDatum(false)); + + return DL_OK; + } + + case BYTEAOID: + { + const char *data; + int32 length; + bytea *result; + + dl_arrow_varlen(column, row, &data, &length); + result = (bytea *) palloc(VARHDRSZ + length); + SET_VARSIZE(result, VARHDRSZ + length); + memcpy(VARDATA(result), data, length); + *value = PointerGetDatum(result); + return DL_OK; + } + + case DATEOID: + { + int32 days = DL_ARROW_VALUES(column, int32)[row]; + DateADT date; + + /* + * Shifting the epoch is a subtraction that can leave the range + * of the type it lands in, so the guard has to come first: by + * the time an overflowed value could be checked it is already + * a different, plausible-looking date. + */ + if (days < DATETIME_MIN_JULIAN - UNIX_EPOCH_JDATE) + return dl_arrow_out_of_range("date"); + + date = days - DL_EPOCH_DELTA_DAYS; + if (!IS_VALID_DATE(date)) + return dl_arrow_out_of_range("date"); + + *value = DateADTGetDatum(date); + return DL_OK; + } + + case TIMESTAMPOID: + case TIMESTAMPTZOID: + { + int64 micros = DL_ARROW_VALUES(column, int64)[row]; + Timestamp ts; + + if (micros < MIN_TIMESTAMP + DL_EPOCH_DELTA_USECS) + return dl_arrow_out_of_range("timestamp"); + + ts = micros - DL_EPOCH_DELTA_USECS; + if (!IS_VALID_TIMESTAMP(ts)) + return dl_arrow_out_of_range("timestamp"); + + *value = TimestampGetDatum(ts); + return DL_OK; + } + + default: + + /* + * Unreachable: dl_arrow_decode_check() refused every type this + * switch does not list. + */ + *isnull = true; + return dl_arrow_decode_refuse(NULL, atttypid); + } +} diff --git a/contrib/datalake_fdw/src/format/arrow_decode.h b/contrib/datalake_fdw/src/format/arrow_decode.h new file mode 100644 index 00000000000..184a1edff2a --- /dev/null +++ b/contrib/datalake_fdw/src/format/arrow_decode.h @@ -0,0 +1,75 @@ +/*------------------------------------------------------------------------- + * + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + * + * arrow_decode.h + * PostgreSQL values out of an Arrow batch. + * + * The read half of the boundary the format layer is built on, and the mirror of + * arrow_builder.h. This side is C: turning a column into Datums means + * allocating text and bytea, an allocation can fail, and a failure in + * PostgreSQL unwinds with longjmp -- which is safe here and would not be if it + * had to pass through C++ frames on the way out. + * + * It reads the buffers of the Arrow C data interface directly rather than + * handing them back to Arrow, which keeps the read path free of C++ and makes + * it a real check on what our own writer exports. + * + * postgres.h must be included before this header. + * + * IDENTIFICATION + * contrib/datalake_fdw/src/format/arrow_decode.h + * + *------------------------------------------------------------------------- + */ + +#ifndef DL_ARROW_DECODE_H +#define DL_ARROW_DECODE_H + +#include "common/dl_err.h" +#include "format/format.h" + +/* + * Whether a column of this Arrow type can be read as this PostgreSQL type. + * Called once per column per batch: the answer depends only on the schema, and + * checking it per value would be the same answer several million times. + * + * `field` is one child of the batch's schema. + */ +extern DlErrCode dl_arrow_decode_check(const struct ArrowSchema *field, + Oid atttypid); + +/* + * One value. Only valid for a column dl_arrow_decode_check() accepted, which + * is what lets this trust the buffer layout instead of re-deriving it. + * + * Values that point at memory -- text, bytea -- are copied into the current + * memory context, because the batch is released long before the tuples built + * from it are done with. + * + * `atttypmod` is the modifier the column was declared with, or -1. A file this + * module did not write has no idea what it was, so a char(n) in it need not be + * padded to n and a varchar(n) need not be within n; without applying it, a + * value that breaks the type's own rules would reach the executor. + */ +extern DlErrCode dl_arrow_decode_value(const struct ArrowArray *column, + int64_t row, Oid atttypid, + int32 atttypmod, + Datum *value, bool *isnull); + +#endif /* DL_ARROW_DECODE_H */ diff --git a/contrib/datalake_fdw/src/format/arrow_support.cpp b/contrib/datalake_fdw/src/format/arrow_support.cpp new file mode 100644 index 00000000000..8addade807f --- /dev/null +++ b/contrib/datalake_fdw/src/format/arrow_support.cpp @@ -0,0 +1,179 @@ +/*------------------------------------------------------------------------- + * + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + * + * arrow_support.cpp + * The type mapping and the error translation shared by the Arrow-facing + * parts of this module. + * + * IDENTIFICATION + * contrib/datalake_fdw/src/format/arrow_support.cpp + * + *------------------------------------------------------------------------- + */ + +/* + * Arrow's headers come first throughout this module: PostgreSQL's c.h defines + * Abs, Min and Max as macros, and a template header has no way to defend + * itself against them. + */ +#include +#include + +#include + +#include "format/arrow_support.h" + +extern "C" +{ +#include "catalog/pg_type.h" +} + +DlErrCode +DlArrowStatus(const arrow::Status &status, const char *operation) +{ + DlErrCode code; + + if (status.ok()) + return DL_OK; + + switch (status.code()) + { + case arrow::StatusCode::IOError: + code = DL_ERR_IO; + break; + case arrow::StatusCode::NotImplemented: + code = DL_ERR_NOT_SUPPORTED; + break; + case arrow::StatusCode::Invalid: + case arrow::StatusCode::TypeError: + case arrow::StatusCode::KeyError: + code = DL_ERR_INVALID_OPTION; + break; + default: + code = DL_ERR_INTERNAL; + break; + } + + dl_error_set(code, operation, arrow::Status::CodeAsString(status.code()).c_str(), + status.message().c_str()); + return code; +} + +std::shared_ptr +DlArrowTypeForPgType(Oid atttypid) +{ + switch (atttypid) + { + case BOOLOID: + return arrow::boolean(); + case INT2OID: + return arrow::int16(); + case INT4OID: + return arrow::int32(); + case INT8OID: + return arrow::int64(); + case FLOAT4OID: + return arrow::float32(); + case FLOAT8OID: + return arrow::float64(); + + /* + * All three of PostgreSQL's string types are one Arrow type: the + * length limit is a constraint PostgreSQL enforces before a value + * reaches us, and Parquet has nowhere to record it. A char(n) + * value arrives already padded, so what is written is what + * PostgreSQL stores. + */ + case TEXTOID: + case VARCHAROID: + case BPCHAROID: + return arrow::utf8(); + + case BYTEAOID: + return arrow::binary(); + case DATEOID: + return arrow::date32(); + + /* + * PostgreSQL keeps both timestamp types in microseconds, so + * microseconds is the unit that loses nothing. timestamptz is a + * point in time held in UTC, which is exactly what an Arrow + * timestamp with a "UTC" zone means; timestamp without time zone + * has no zone, and Arrow says that by leaving it empty. + */ + case TIMESTAMPOID: + return arrow::timestamp(arrow::TimeUnit::MICRO); + case TIMESTAMPTZOID: + return arrow::timestamp(arrow::TimeUnit::MICRO, "UTC"); + + default: + return nullptr; + } +} + +std::shared_ptr +DlArrowSchemaFromTupleDesc(TupleDesc tupdesc) +{ + std::vector> fields; + + fields.reserve(tupdesc->natts); + + for (int i = 0; i < tupdesc->natts; i++) + { + Form_pg_attribute attr = TupleDescAttr(tupdesc, i); + + /* + * A dropped column has no type to write and no name worth recording. + * Leaving a placeholder in the file would keep column positions + * aligned, but nothing reads such a file yet, so refusing is the + * answer that cannot be silently wrong. + */ + if (attr->attisdropped) + { + dl_error_set(DL_ERR_NOT_SUPPORTED, "arrow schema", nullptr, + "a dropped column cannot be written to a data file"); + return nullptr; + } + + std::shared_ptr type = DlArrowTypeForPgType(attr->atttypid); + + if (type == nullptr) + { + std::string message = std::string("column \"") + + NameStr(attr->attname) + "\" has a type that lake tables " + "cannot store yet"; + + dl_error_set(DL_ERR_NOT_SUPPORTED, "arrow schema", nullptr, + message.c_str()); + return nullptr; + } + + /* + * Every field is nullable, including one PostgreSQL marked NOT NULL. + * Recording it as required would buy nothing -- PostgreSQL has already + * rejected the nulls before a tuple reaches this layer -- and would + * turn any later relaxation of the constraint into a write failure + * against files already on disk. + */ + fields.push_back(arrow::field(NameStr(attr->attname), type, + /* nullable */ true)); + } + + return arrow::schema(fields); +} diff --git a/contrib/datalake_fdw/src/format/arrow_support.h b/contrib/datalake_fdw/src/format/arrow_support.h new file mode 100644 index 00000000000..d79b5934ddc --- /dev/null +++ b/contrib/datalake_fdw/src/format/arrow_support.h @@ -0,0 +1,75 @@ +/*------------------------------------------------------------------------- + * + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + * + * arrow_support.h + * What every Arrow-facing translation unit in this module needs: how a + * PostgreSQL column type is stored, and how an Arrow failure is reported. + * + * The type mapping is in one place because the writer, the builder that feeds + * it and the reader that decodes what comes back all have to agree on it, and a + * disagreement between them would show up as wrong values rather than as an + * error. + * + * IDENTIFICATION + * contrib/datalake_fdw/src/format/arrow_support.h + * + *------------------------------------------------------------------------- + */ + +#ifndef DL_ARROW_SUPPORT_H +#define DL_ARROW_SUPPORT_H + +#include + +#include +#include +#include + +#include "common/dl_err.h" +#include "common/dl_pg_api.h" + +extern "C" +{ +#include "access/tupdesc.h" +} + +/* + * Turns an Arrow status into this module's error code, recording what Arrow + * said -- its own class and message are the only thing that makes a failure in + * a third-party library diagnosable, and the code alone throws them away. + * `operation` names what was being attempted. A successful status records + * nothing and returns DL_OK, so call sites can wrap every Arrow call. + */ +extern DlErrCode DlArrowStatus(const arrow::Status &status, const char *operation); + +/* + * The Arrow type a column of this PostgreSQL type is stored as, or a null + * pointer when the type has no mapping yet. Callers report the refusal + * themselves, because only they know which column it was about. + */ +extern std::shared_ptr DlArrowTypeForPgType(Oid atttypid); + +/* + * The whole descriptor. Returns a null pointer and records which column was + * the problem in the error detail: a type with no mapping, or a dropped column, + * which has no type to write and which nothing reads a file for yet. + */ +extern std::shared_ptr DlArrowSchemaFromTupleDesc(TupleDesc tupdesc); + +#endif /* DL_ARROW_SUPPORT_H */ diff --git a/contrib/datalake_fdw/src/format/format.h b/contrib/datalake_fdw/src/format/format.h index 003101ae550..db793eb9689 100644 --- a/contrib/datalake_fdw/src/format/format.h +++ b/contrib/datalake_fdw/src/format/format.h @@ -64,10 +64,44 @@ struct ArrowArray { }; #endif /* ARROW_C_DATA_INTERFACE */ -typedef struct Fragment Fragment; /* opaque in skeleton */ -typedef struct ProjectionSet ProjectionSet; +/* + * One unit of read work. A fragment is a range of row groups rather than a + * whole file, because that is the granularity a scan can be divided at: several + * segments can then read one large file at once, which file-at-a-time + * assignment cannot express. + */ +typedef struct Fragment +{ + const char *path; + int first_row_group; /* 0-based */ + int n_row_groups; /* 0 == to the end of the file */ +} Fragment; + +/* + * The columns to materialise, as 0-based indexes into the file schema. A NULL + * set, or one with no columns, means every column: "read nothing" is not a + * projection anyone asks for, so it is not worth a second way to say "all". + */ +typedef struct ProjectionSet +{ + const int *columns; + int ncolumns; +} ProjectionSet; + +/* + * A writer holds a whole row group before it can write one, so this is a bound + * on memory as much as on the file's shape. It is Parquet's own default + * maximum, which is what makes it a sane ceiling for any format. + */ +#define DL_MAX_ROW_GROUP_ROWS (1024 * 1024) + +typedef struct WriterOptions +{ + const char *compression; /* format-defined name; NULL for the default */ + int64_t row_group_size; /* rows per row group; 0 for the default */ +} WriterOptions; + typedef struct RowGroupFilterSet RowGroupFilterSet; -typedef struct WriterOptions WriterOptions; typedef struct FileMeta FileMeta; typedef struct DeleteFileSet DeleteFileSet; @@ -75,28 +109,49 @@ typedef struct DeleteFileSet DeleteFileSet; * No global slots or trampolines, ever. */ typedef struct FormatReader FormatReader; typedef struct FormatReaderOps { - /* Each batch yields ArrowArray+ArrowSchema; last column is a hidden int64 file-row - * ordinal (for MoR positional deletes). */ + /* Each batch yields ArrowArray+ArrowSchema. A hidden trailing int64 column + * carrying the file-row ordinal is what merge-on-read positional deletes + * will match against; it arrives with them, so a batch is the projected + * columns and nothing else for now. */ DlErrCode (*next_batch)(FormatReader *, struct ArrowArray *out, struct ArrowSchema *schema, bool *eof); - void (*close)(FormatReader *); /* void cleanup ABI: noexcept, idempotent, never ereport */ + /* void cleanup ABI: noexcept, never ereport. Takes the caller's handle so + * that it can clear it -- these run on the resource-owner path during + * abort, where the same cleanup can be reached twice, and a second call + * has to find nothing left rather than a freed reader. */ + void (*close)(FormatReader **); } FormatReaderOps; struct FormatReader { const FormatReaderOps *ops; void *impl; }; typedef struct FormatWriter FormatWriter; typedef struct FormatWriterOps { - DlErrCode (*write_batch)(FormatWriter *, struct ArrowArray *batch); /* success == consumed */ + /* + * The batch is consumed whether or not the write succeeds: an + * implementation hands it to a library that takes ownership at the call, + * and there is no point at which it could hand it back. The caller is + * left with a released ArrowArray either way. + */ + DlErrCode (*write_batch)(FormatWriter *, struct ArrowArray *batch); /* Rolling support: actual bytes encoded into the sink so far. Valid to query after a * successful write_batch; on failure returns an error code and *out is invalid. * The write.c orchestration layer rolls files (finish -> new open_writer) when this - * reaches the soft target; overshoot of at most one batch is allowed. */ + * reaches the soft target. A format writes in units it cannot split -- a Parquet + * row group is one -- and what has not been written is not counted, so the target + * is overshot by at most one such unit. */ DlErrCode (*bytes_written)(FormatWriter *, int64_t *out); - DlErrCode (*finish)(FormatWriter *, FileMeta **meta); /* reportable close-time errors - * surface ONLY here */ - void (*abort)(FormatWriter *); /* void cleanup ABI: noexcept, idempotent, never ereport */ + /* Reportable close-time errors surface ONLY here. The writer is consumed + * and the caller's handle cleared whether or not it succeeds: a file whose + * footer could not be written is not one anything can retry against. */ + DlErrCode (*finish)(FormatWriter **, FileMeta **meta); + /* void cleanup ABI, as for close() above: discards the file being written + * and clears the caller's handle. */ + void (*abort)(FormatWriter **); } FormatWriterOps; struct FormatWriter { const FormatWriterOps *ops; void *impl; }; +/* Bumped when an existing field changes meaning; appending does not need it. */ +#define DL_FORMAT_ABI_VERSION 1 + typedef struct FormatRoutine { uint32_t abi_version, struct_size; /* same prefix-compat semantics as meta engine */ const char *name; /* "parquet" */ diff --git a/contrib/datalake_fdw/src/format/format_registry.c b/contrib/datalake_fdw/src/format/format_registry.c index 83370a51f64..44bf9ec0732 100644 --- a/contrib/datalake_fdw/src/format/format_registry.c +++ b/contrib/datalake_fdw/src/format/format_registry.c @@ -26,15 +26,35 @@ *------------------------------------------------------------------------- */ +#include #include +#include +#include "common/dl_err.h" #include "format/format.h" +#include "format/parquet/parquet_format.h" -/* No formats in the skeleton; parquet lands in PR-3/4. Callers must treat - * NULL as not-supported. */ +/* + * Parquet is the only format so far. A name that reaches here came from a + * table option, so an unknown one is an ordinary mistake and the caller has to + * be able to say which name it was -- returning a bare NULL would leave every + * caller to write that message again, and get it wrong differently. The name + * goes into the error detail, so a caller that reports DL_ERR_NOT_SUPPORTED + * gets it without knowing this function exists. + */ const FormatRoutine * GetFormatRoutine(const char *format) { + char message[128]; + + if (format != NULL && strcmp(format, "parquet") == 0) + return GetParquetFormatRoutine(); + + snprintf(message, sizeof(message), + "\"%s\" is not a data file format this build can read or write", + format == NULL ? "" : format); + dl_error_set(DL_ERR_NOT_SUPPORTED, "get_format", NULL, message); + return NULL; } diff --git a/contrib/datalake_fdw/src/format/parquet/parquet_format.cpp b/contrib/datalake_fdw/src/format/parquet/parquet_format.cpp new file mode 100644 index 00000000000..ac094389ee9 --- /dev/null +++ b/contrib/datalake_fdw/src/format/parquet/parquet_format.cpp @@ -0,0 +1,55 @@ +/*------------------------------------------------------------------------- + * + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + * + * parquet_format.cpp + * Parquet as a format this module can read and write. + * + * Parquet is reached through Arrow rather than through libparquet on its own, + * because libparquet is written in terms of Arrow's types: linking it already + * links Arrow, and going around Arrow would mean re-deriving the definition and + * repetition levels, the four ways a decimal can be stored, and the timestamp + * unit rules that arrow::parquet already gets right. + * + * Nothing here reads or writes anything but a local file yet. The storage + * facade in common/file_system_wrapper.h is where object storage arrives, as an + * arrow::io::RandomAccessFile over it; parquet_read.cpp and parquet_write.cpp + * are the only files that have to change when it does. + * + * IDENTIFICATION + * contrib/datalake_fdw/src/format/parquet/parquet_format.cpp + * + *------------------------------------------------------------------------- + */ + +#include "format/parquet/parquet_format.h" +#include "format/parquet/parquet_internal.h" + +static const FormatRoutine parquet_format_routine = { + DL_FORMAT_ABI_VERSION, + sizeof(FormatRoutine), + "parquet", + parquet_open_reader, + parquet_open_writer +}; + +extern "C" const FormatRoutine * +GetParquetFormatRoutine(void) +{ + return &parquet_format_routine; +} diff --git a/contrib/datalake_fdw/src/format/parquet/parquet_format.h b/contrib/datalake_fdw/src/format/parquet/parquet_format.h new file mode 100644 index 00000000000..20f3ff788fa --- /dev/null +++ b/contrib/datalake_fdw/src/format/parquet/parquet_format.h @@ -0,0 +1,45 @@ +/*------------------------------------------------------------------------- + * + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + * + * parquet_format.h + * The Parquet reader and writer. + * + * IDENTIFICATION + * contrib/datalake_fdw/src/format/parquet/parquet_format.h + * + *------------------------------------------------------------------------- + */ + +#ifndef DL_PARQUET_FORMAT_H +#define DL_PARQUET_FORMAT_H + +#include "format/format.h" + +#ifdef __cplusplus +extern "C" +{ +#endif + +extern const FormatRoutine *GetParquetFormatRoutine(void); + +#ifdef __cplusplus +} +#endif + +#endif /* DL_PARQUET_FORMAT_H */ diff --git a/contrib/datalake_fdw/src/format/parquet/parquet_internal.h b/contrib/datalake_fdw/src/format/parquet/parquet_internal.h new file mode 100644 index 00000000000..f7acbdd2f94 --- /dev/null +++ b/contrib/datalake_fdw/src/format/parquet/parquet_internal.h @@ -0,0 +1,47 @@ +/*------------------------------------------------------------------------- + * + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + * + * parquet_internal.h + * What the halves of the Parquet format say to each other. + * + * Reading and writing a Parquet file have nothing in common but the name of + * the format, so they are separate translation units; this is the only thing + * they share, and parquet_format.cpp is the only other file that needs it. + * + * IDENTIFICATION + * contrib/datalake_fdw/src/format/parquet/parquet_internal.h + * + *------------------------------------------------------------------------- + */ + +#ifndef DL_PARQUET_INTERNAL_H +#define DL_PARQUET_INTERNAL_H + +#include "format/format.h" + +extern DlErrCode parquet_open_reader(const Fragment *fragment, + const ProjectionSet *projection, + const RowGroupFilterSet *filters, + FormatReader **out); + +extern DlErrCode parquet_open_writer(const char *path, void *tupdesc, + const WriterOptions *options, + FormatWriter **out); + +#endif /* DL_PARQUET_INTERNAL_H */ diff --git a/contrib/datalake_fdw/src/format/parquet/parquet_read.cpp b/contrib/datalake_fdw/src/format/parquet/parquet_read.cpp new file mode 100644 index 00000000000..9c245672fbe --- /dev/null +++ b/contrib/datalake_fdw/src/format/parquet/parquet_read.cpp @@ -0,0 +1,280 @@ +/*------------------------------------------------------------------------- + * + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + * + * parquet_read.cpp + * Reading a Parquet file. + * + * IDENTIFICATION + * contrib/datalake_fdw/src/format/parquet/parquet_read.cpp + * + *------------------------------------------------------------------------- + */ + +#include +#include +#include + +#include +#include +#include +#include + +#include "format/arrow_support.h" + +#include "am_iceberg/pg_iceberg_guc.h" +#include "common/dl_resource.h" +#include "common/dl_wrappers.h" +#include "format/parquet/parquet_internal.h" + +struct ParquetReader +{ + FormatReader base; + std::shared_ptr file; + std::unique_ptr reader; + std::shared_ptr batches; +}; + +static DlErrCode +parquet_reader_next_batch(FormatReader *reader, struct ArrowArray *out, + struct ArrowSchema *schema, bool *eof) +{ + DlErrCode result = DL_OK; + + if (reader == NULL || out == NULL || eof == NULL) + return DL_ARG_ERROR("next_batch"); + + *eof = false; + + DL_ABI_GUARD_BEGIN + { + ParquetReader *impl = static_cast(reader->impl); + std::shared_ptr batch; + arrow::Status status = impl->batches->ReadNext(&batch); + + if (!status.ok()) + return DlArrowStatus(status, "read a Parquet batch"); + + if (batch == nullptr) + { + *eof = true; + return DL_OK; + } + + return DlArrowStatus(arrow::ExportRecordBatch(*batch, out, schema), + "export a Parquet batch"); + } + DL_ABI_GUARD_END(result, "next_batch"); + + return result; +} + +/* + * What the resource owner calls if nothing else did. C linkage because a C + * function pointer is what it is handed to. + */ +extern "C" void +parquet_reader_release(void *arg) +{ + DL_CLEANUP_GUARD_BEGIN + { + delete static_cast(arg); + } + DL_CLEANUP_GUARD_END; +} + +static void +parquet_reader_close(FormatReader **reader) +{ + if (reader == NULL || *reader == NULL) + return; + + /* + * Clear the caller's handle first. DL_CLEANUP_GUARD_END can elog(WARNING), + * and an escalation there would longjmp past the assignment -- leaving the + * caller holding a reader that has already been released, which is the + * thing taking the handle by address exists to prevent. + */ + ParquetReader *impl = static_cast((*reader)->impl); + + *reader = NULL; + dl_resource_forget(parquet_reader_release, impl); + + parquet_reader_release(impl); +} + +static const FormatReaderOps parquet_reader_ops = { + parquet_reader_next_batch, + parquet_reader_close +}; + +/* + * Which row groups this fragment covers. A fragment is a range rather than a + * whole file so that one large file can be read by several segments at once; + * an empty range is legal and reads nothing. + */ +static DlErrCode +parquet_row_groups(const Fragment *fragment, int total, + std::vector *row_groups) +{ + int first = fragment->first_row_group; + int count = fragment->n_row_groups; + + /* + * The last test is written as a subtraction because the addition it + * replaces overflows: first + INT_MAX wraps negative, passes the check, and + * the loop below then builds a two-billion-element vector out of a range + * that should have been rejected. + */ + if (first < 0 || first > total || count < 0 || count > total - first) + { + char message[160]; + + snprintf(message, sizeof(message), + "row groups %d..%d were asked for from a file that has %d", + first, count > 0 ? first + count - 1 : first, total); + dl_error_set(DL_ERR_INVALID_OPTION, "open a Parquet file", NULL, message); + return DL_ERR_INVALID_OPTION; + } + + if (count == 0) + count = total - first; + + for (int i = 0; i < count; i++) + row_groups->push_back(first + i); + + return DL_OK; +} + +DlErrCode +parquet_open_reader(const Fragment *fragment, const ProjectionSet *projection, + const RowGroupFilterSet *filters, FormatReader **out) +{ + DlErrCode result = DL_OK; + + if (out == NULL) + return DL_ARG_ERROR("open_reader"); + *out = NULL; + + if (fragment == NULL || fragment->path == NULL) + return DL_ARG_ERROR("open_reader"); + + /* + * Statistics-based row group pruning is not implemented. Accepting the + * filters and ignoring them would still give the right rows, so nothing + * would fail -- which is exactly why it is refused instead: a caller that + * believed the pruning had happened would have no way to find out. + */ + if (filters != NULL) + { + dl_error_set(DL_ERR_NOT_SUPPORTED, "open a Parquet file", NULL, + "row group filtering is not implemented yet"); + return DL_ERR_NOT_SUPPORTED; + } + + DL_ABI_GUARD_BEGIN + { + std::unique_ptr impl(new ParquetReader()); + arrow::MemoryPool *pool = arrow::default_memory_pool(); + std::vector row_groups; + std::vector columns; + DlErrCode rc; + + parquet::arrow::FileReaderBuilder builder; + parquet::ArrowReaderProperties properties; + + arrow::Result> file = + arrow::io::ReadableFile::Open(fragment->path, pool); + + if (!file.ok()) + return DlArrowStatus(file.status(), "open a Parquet file"); + impl->file = *file; + + arrow::Status status = builder.Open(impl->file); + + if (!status.ok()) + return DlArrowStatus(status, "open a Parquet file"); + + /* The same batch size the write side accumulates to. */ + properties.set_batch_size(iceberg_batch_rows); + + /* + * A backend is not a thread pool. Arrow will read column chunks in + * parallel if asked, and a worker thread that hits an error has no way + * to report it through PostgreSQL's error handling, so this reads on + * the thread it was called on. + */ + properties.set_use_threads(false); + + status = builder.memory_pool(pool)->properties(properties) + ->Build(&impl->reader); + + if (!status.ok()) + return DlArrowStatus(status, "open a Parquet file"); + + rc = parquet_row_groups(fragment, impl->reader->num_row_groups(), + &row_groups); + if (rc != DL_OK) + return rc; + + if (projection != NULL && projection->ncolumns > 0) + columns.assign(projection->columns, + projection->columns + projection->ncolumns); + else + { + std::shared_ptr schema; + + status = impl->reader->GetSchema(&schema); + if (!status.ok()) + return DlArrowStatus(status, "read a Parquet schema"); + + for (int i = 0; i < schema->num_fields(); i++) + columns.push_back(i); + } + + /* + * Arrow 21 deprecates this in favour of a Result-returning one that + * Arrow 9 does not have, so whoever raises the floor past 21 gets a + * warning here and a version guard to write -- the one around + * FileWriter::Open in parquet_write.cpp is the shape of it. + */ + status = impl->reader->GetRecordBatchReader(row_groups, columns, + &impl->batches); + if (!status.ok()) + return DlArrowStatus(status, "open a Parquet batch reader"); + + impl->base.ops = &parquet_reader_ops; + impl->base.impl = impl.get(); + + /* + * The file is open from here, so this is the last thing that may fail: + * past it, nothing can lose track of the descriptor. + */ + if (!dl_resource_remember(parquet_reader_release, impl.get())) + { + dl_error_set(DL_ERR_INTERNAL, "open_reader", NULL, + "could not record the open file for cleanup"); + return DL_ERR_INTERNAL; + } + + *out = &impl.release()->base; + } + DL_ABI_GUARD_END(result, "open_reader"); + + return result; +} diff --git a/contrib/datalake_fdw/src/format/parquet/parquet_write.cpp b/contrib/datalake_fdw/src/format/parquet/parquet_write.cpp new file mode 100644 index 00000000000..e5c00b9b043 --- /dev/null +++ b/contrib/datalake_fdw/src/format/parquet/parquet_write.cpp @@ -0,0 +1,451 @@ +/*------------------------------------------------------------------------- + * + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + * + * parquet_write.cpp + * Writing a Parquet file. + * + * IDENTIFICATION + * contrib/datalake_fdw/src/format/parquet/parquet_write.cpp + * + *------------------------------------------------------------------------- + */ + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "format/arrow_support.h" + +#include "common/dl_resource.h" +#include "common/dl_wrappers.h" +#include "format/parquet/parquet_internal.h" + +struct ParquetWriter +{ + FormatWriter base; + std::string path; + std::shared_ptr schema; + std::shared_ptr sink; + std::unique_ptr writer; + + /* + * A row group is written whole, so the batches that go into one are held + * until there are enough of them. This is not a buffer we chose to add: + * Parquet cannot begin a row group it does not have, and the alternative + * -- one row group per batch -- would produce files whose row groups are + * a thousand rows, where a reader expects something nearer a million and + * pays a seek for each one. + */ + std::vector> pending; + int64_t pending_rows; + int64_t row_group_size; +}; + +/* + * Releases a batch unless something already has. The C data interface clears + * the callback when ownership moves, so this is a no-op on the path where Arrow + * took the batch, and the release on every other path. + */ +class ParquetReleaseBatch +{ +public: + explicit ParquetReleaseBatch(struct ArrowArray *batch) : batch_(batch) {} + ~ParquetReleaseBatch() + { + if (batch_ != nullptr && batch_->release != nullptr) + batch_->release(batch_); + } + +private: + struct ArrowArray *batch_; +}; + +/* + * Gives up on the file being written. The sink is closed first and the writer + * left to its destructor: a Parquet writer writes the footer when it closes, + * and against a sink that is already closed it cannot -- which is what stops a + * complete, valid, truncated file appearing at the path if the unlink does not + * take. What it leaves then has no footer, so nothing can read it, and that is + * why the unlink's result is not worth reporting. + */ +static void +parquet_discard(ParquetWriter *impl) +{ + if (impl->sink != nullptr) + (void) impl->sink->Close(); + + (void) unlink(impl->path.c_str()); +} + +/* + * What the resource owner calls if nothing else did. A file that got this far + * was never finished, so it is discarded rather than left: the same thing + * abort() does, and for the same reason. C linkage because a C function + * pointer is what it is handed to. + */ +extern "C" void +parquet_writer_release(void *arg) +{ + DL_CLEANUP_GUARD_BEGIN + { + std::unique_ptr impl(static_cast(arg)); + + parquet_discard(impl.get()); + } + DL_CLEANUP_GUARD_END; +} + +/* + * Discards the file unless the scope it guards clears it. finish() has to get + * rid of a file it could not complete on every way out, and one of those ways + * is an exception that the guard around it turns into an error code -- past any + * cleanup written as a statement. + */ +class ParquetDiscardOnFailure +{ +public: + explicit ParquetDiscardOnFailure(ParquetWriter *writer) : writer_(writer) {} + ~ParquetDiscardOnFailure() + { + if (writer_ != nullptr) + parquet_discard(writer_); + } + void Keep() { writer_ = nullptr; } + +private: + ParquetWriter *writer_; +}; + +static arrow::Status +parquet_flush_row_group(ParquetWriter *impl) +{ + if (impl->pending.empty()) + return arrow::Status::OK(); + + ARROW_ASSIGN_OR_RAISE(std::shared_ptr table, + arrow::Table::FromRecordBatches(impl->schema, + impl->pending)); + + impl->pending.clear(); + impl->pending_rows = 0; + + return impl->writer->WriteTable(*table, impl->row_group_size); +} + +static DlErrCode +parquet_writer_write_batch(FormatWriter *writer, struct ArrowArray *batch) +{ + DlErrCode result = DL_OK; + + if (batch == NULL) + return DL_ARG_ERROR("write_batch"); + + /* + * The interface promises the batch is consumed whether or not the write + * succeeds, and that has to hold for the ways out that are not a return: + * importing allocates before it takes ownership, so it can throw with the + * batch still live. On the ordinary path the import has already cleared + * the callback and this does nothing. + */ + ParquetReleaseBatch release_batch(batch); + + if (writer == NULL) + return DL_ARG_ERROR("write_batch"); + + DL_ABI_GUARD_BEGIN + { + ParquetWriter *impl = static_cast(writer->impl); + + arrow::Result> imported = + arrow::ImportRecordBatch(batch, impl->schema); + + if (!imported.ok()) + return DlArrowStatus(imported.status(), "import an Arrow batch"); + + int64_t rows = (*imported)->num_rows(); + + /* + * Flush before the batch that would take the group past its size, not + * after. Flushing afterwards leaves a remainder that WriteTable emits + * as a second, tiny row group -- so a batch size that does not divide + * the row group size would produce exactly the file of many small row + * groups this buffering exists to avoid. + */ + if (impl->pending_rows > 0 && + impl->pending_rows + rows > impl->row_group_size) + { + arrow::Status status = parquet_flush_row_group(impl); + + if (!status.ok()) + return DlArrowStatus(status, "write a Parquet row group"); + } + + impl->pending_rows += rows; + impl->pending.push_back(*imported); + } + DL_ABI_GUARD_END(result, "write_batch"); + + return result; +} + +static DlErrCode +parquet_writer_bytes_written(FormatWriter *writer, int64_t *out) +{ + DlErrCode result = DL_OK; + + if (writer == NULL || out == NULL) + return DL_ARG_ERROR("bytes_written"); + + DL_ABI_GUARD_BEGIN + { + ParquetWriter *impl = static_cast(writer->impl); + + /* + * What has reached the file, which trails what has been handed over: + * a row group is written whole, so the batches waiting for one are + * not in this number. The layer that rolls files reads it to decide + * when a file is big enough, and the undercount costs it one row group + * of overshoot -- the tolerance the interface is written with. + */ + arrow::Result position = impl->sink->Tell(); + + if (!position.ok()) + return DlArrowStatus(position.status(), "measure a Parquet file"); + + *out = *position; + } + DL_ABI_GUARD_END(result, "bytes_written"); + + return result; +} + +static DlErrCode +parquet_writer_finish(FormatWriter **writer, FileMeta **meta) +{ + DlErrCode result = DL_OK; + + if (writer == NULL || *writer == NULL) + return DL_ARG_ERROR("finish_writer"); + + if (meta != NULL) + *meta = NULL; /* what a commit needs is not collected yet */ + + DL_ABI_GUARD_BEGIN + { + /* + * Taking ownership here is what makes "consumed either way" true even + * of the paths that leave through an exception: the caller's handle is + * cleared before anything that could fail. + */ + std::unique_ptr impl( + static_cast((*writer)->impl)); + ParquetDiscardOnFailure discard(impl.get()); + + *writer = NULL; + dl_resource_forget(parquet_writer_release, impl.get()); + + /* + * Closing the writer is what writes the footer, so a failure here + * leaves an unreadable file behind and has to be reported -- this is + * the one place in the writer's interface where close-time errors can + * still reach a caller. + */ + arrow::Status status = parquet_flush_row_group(impl.get()); + + if (status.ok()) + status = impl->writer->Close(); + if (status.ok()) + status = impl->sink->Close(); + + if (!status.ok()) + return DlArrowStatus(status, "finish a Parquet file"); + + discard.Keep(); + } + DL_ABI_GUARD_END(result, "finish_writer"); + + return result; +} + +static void +parquet_writer_abort(FormatWriter **writer) +{ + if (writer == NULL || *writer == NULL) + return; + + /* Cleared first; see the note in parquet_reader_close(). */ + ParquetWriter *impl = static_cast((*writer)->impl); + + *writer = NULL; + dl_resource_forget(parquet_writer_release, impl); + + /* + * A file that was never finished has no footer, so nothing can read it and + * leaving it behind only costs space and confusion. Failures are ignored: + * this runs while an error is already being handled. + */ + parquet_writer_release(impl); +} + +static const FormatWriterOps parquet_writer_ops = { + parquet_writer_write_batch, + parquet_writer_bytes_written, + parquet_writer_finish, + parquet_writer_abort +}; + +static DlErrCode +parquet_compression(const char *name, arrow::Compression::type *out) +{ + std::string requested(name); + char message[128]; + + if (requested == "none" || requested == "uncompressed") + *out = arrow::Compression::UNCOMPRESSED; + else if (requested == "snappy") + *out = arrow::Compression::SNAPPY; + else if (requested == "gzip") + *out = arrow::Compression::GZIP; + else if (requested == "zstd") + *out = arrow::Compression::ZSTD; + else + { + snprintf(message, sizeof(message), + "\"%s\" is not a compression this build can write", name); + dl_error_set(DL_ERR_INVALID_OPTION, "open a Parquet file", NULL, message); + return DL_ERR_INVALID_OPTION; + } + + return DL_OK; +} + +DlErrCode +parquet_open_writer(const char *path, void *tupdesc_arg, + const WriterOptions *options, FormatWriter **out) +{ + DlErrCode result = DL_OK; + + if (out == NULL) + return DL_ARG_ERROR("open_writer"); + *out = NULL; + + if (path == NULL || tupdesc_arg == NULL) + return DL_ARG_ERROR("open_writer"); + + DL_ABI_GUARD_BEGIN + { + std::unique_ptr impl(new ParquetWriter()); + arrow::MemoryPool *pool = arrow::default_memory_pool(); + parquet::WriterProperties::Builder properties; + arrow::Compression::type compression = arrow::Compression::SNAPPY; + DlErrCode rc; + + impl->path = path; + impl->schema = DlArrowSchemaFromTupleDesc((TupleDesc) tupdesc_arg); + if (impl->schema == nullptr) + return DL_ERR_NOT_SUPPORTED; /* detail already recorded */ + + if (options != NULL && options->compression != NULL) + { + rc = parquet_compression(options->compression, &compression); + if (rc != DL_OK) + return rc; + } + properties.compression(compression); + + if (options != NULL && options->row_group_size > 0) + properties.max_row_group_length(options->row_group_size); + + std::shared_ptr built = properties.build(); + + /* Whether it was asked for or left to Parquet, this is the size. */ + impl->row_group_size = built->max_row_group_length(); + impl->pending_rows = 0; + + arrow::Result> sink = + arrow::io::FileOutputStream::Open(path); + + if (!sink.ok()) + return DlArrowStatus(sink.status(), "create a Parquet file"); + impl->sink = *sink; + + /* + * The Arrow schema is deliberately not stored in the file's metadata. + * With it, reading back would restore the types from our own note + * rather than from Parquet's, and a round trip would agree with itself + * no matter what it had written; without it, what comes back is what + * any other reader of the file sees. + */ + /* + * Arrow 11 deprecated the form that returns its writer through an out + * parameter. Both spellings have to be here because the versions this + * builds against range from 9 to 17 depending on the distribution, and + * the older one warns on the newer Arrow rather than failing -- which is + * the kind of warning that stops being read. + */ +#if ARROW_VERSION_MAJOR >= 11 + arrow::Result> writer = + parquet::arrow::FileWriter::Open(*impl->schema, pool, impl->sink, + built, + parquet::default_arrow_writer_properties()); + arrow::Status status = writer.status(); + + if (status.ok()) + impl->writer = std::move(*writer); +#else + arrow::Status status = + parquet::arrow::FileWriter::Open(*impl->schema, pool, impl->sink, + built, + parquet::default_arrow_writer_properties(), + &impl->writer); +#endif + + if (!status.ok()) + { + (void) impl->sink->Close(); + unlink(path); + return DlArrowStatus(status, "create a Parquet file"); + } + + impl->base.ops = &parquet_writer_ops; + impl->base.impl = impl.get(); + + /* Last thing that may fail; see parquet_open_reader(). */ + if (!dl_resource_remember(parquet_writer_release, impl.get())) + { + parquet_discard(impl.get()); + dl_error_set(DL_ERR_INTERNAL, "open_writer", NULL, + "could not record the open file for cleanup"); + return DL_ERR_INTERNAL; + } + + *out = &impl.release()->base; + } + DL_ABI_GUARD_END(result, "open_writer"); + + return result; +} diff --git a/contrib/datalake_fdw/src/test/datalake_fdw_test.c b/contrib/datalake_fdw/src/test/datalake_fdw_test.c new file mode 100644 index 00000000000..1d2c2946d11 --- /dev/null +++ b/contrib/datalake_fdw/src/test/datalake_fdw_test.c @@ -0,0 +1,411 @@ +/*------------------------------------------------------------------------- + * + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + * + * datalake_fdw_test.c + * The format layer, reachable from SQL. + * + * A data file is written and read by the access method, which is not finished; + * until it is, there is no way to run the format layer in a real backend, and + * "it compiles" would be the only thing anyone could say about it. These two + * functions are that way in: they write the result of a query to a file and + * read a file back as rows, so a round trip is an ordinary SQL statement. + * + * They are a separate extension because they are not part of what this module + * offers -- installing datalake_fdw does not put them in the database. + * + * IDENTIFICATION + * contrib/datalake_fdw/src/test/datalake_fdw_test.c + * + *------------------------------------------------------------------------- + */ + +#include "postgres.h" + +#include "executor/spi.h" +#include "funcapi.h" +#include "utils/builtins.h" +#include "utils/memutils.h" +#include "utils/tuplestore.h" + +#include "am_iceberg/pg_iceberg_guc.h" +#include "common/dl_err.h" +#include "format/arrow_builder.h" +#include "format/arrow_decode.h" +#include "format/format.h" + +PG_FUNCTION_INFO_V1(datalake_parquet_write); +PG_FUNCTION_INFO_V1(datalake_parquet_read); + +static const FormatRoutine * +parquet_routine(void) +{ + const FormatRoutine *routine = GetFormatRoutine("parquet"); + + /* + * Only reachable from a build that dropped the format, so what it can say + * is whatever the registry recorded -- guessing at a reason here would be + * a message that outlives the thing it describes. + */ + if (routine == NULL) + dl_error_report(ERROR, DL_ERR_NOT_SUPPORTED, "get_format"); + + return routine; +} + +/* + * Hands one batch to the writer. The batch is consumed either way, so there is + * nothing left to release when this reports a failure. + */ +static void +write_one_batch(FormatWriter *writer, DlArrowBuilder builder) +{ + struct ArrowArray batch; + DlErrCode rc; + + rc = dl_arrow_builder_flush(builder, &batch); + if (rc != DL_OK) + dl_error_report(ERROR, rc, "build_batch"); + + rc = writer->ops->write_batch(writer, &batch); + if (rc != DL_OK) + dl_error_report(ERROR, rc, "write_batch"); +} + +/* + * datalake_parquet_write(path, query, row_group_size) -> rows written + * + * The rows the query returns are written to `path` as Parquet. A row group + * size of zero leaves the format's own default in place; anything else also + * becomes the number of rows per batch, because a row group is closed at a + * batch boundary and the option would otherwise be rounded away by a batch size + * that does not divide by it. + */ +Datum +datalake_parquet_write(PG_FUNCTION_ARGS) +{ + char *path = text_to_cstring(PG_GETARG_TEXT_PP(0)); + char *query = text_to_cstring(PG_GETARG_TEXT_PP(1)); + int32 row_group_size = PG_GETARG_INT32(2); + const FormatRoutine *routine = parquet_routine(); + WriterOptions options = {0}; + FormatWriter *volatile open_writer = NULL; + DlArrowBuilder volatile open_builder = NULL; + long batch_rows = iceberg_batch_rows; + int64 written = 0; + MemoryContext row_context; + + /* + * Bounded above as well as below, and by the same number as + * iceberg.batch_rows: a row group is held in memory until it is complete, + * so an unbounded one asks the writer to buffer the whole result set. + */ + if (row_group_size < 0 || row_group_size > DL_MAX_ROW_GROUP_ROWS) + ereport(ERROR, + (errcode(ERRCODE_INVALID_PARAMETER_VALUE), + errmsg("row group size must be between 0 and %d", + DL_MAX_ROW_GROUP_ROWS))); + + options.row_group_size = row_group_size; + if (row_group_size > 0 && row_group_size < batch_rows) + batch_rows = row_group_size; + + if (SPI_connect() != SPI_OK_CONNECT) + elog(ERROR, "SPI_connect failed"); + + /* + * Detoasting a value allocates, and the copy is dead as soon as it has been + * appended. Without a context of its own, a wide table would hold every + * copy it ever made until the function returned. + */ + row_context = AllocSetContextCreate(CurrentMemoryContext, + "datalake_parquet_write", + ALLOCSET_DEFAULT_SIZES); + + PG_TRY(); + { + SPIPlanPtr plan; + Portal portal; + TupleDesc tupdesc = NULL; + FormatWriter *writer = NULL; + DlArrowBuilder builder = NULL; + Datum *values = NULL; + bool *nulls = NULL; + DlErrCode rc; + + plan = SPI_prepare(query, 0, NULL); + if (plan == NULL) + elog(ERROR, "SPI_prepare failed: %s", + SPI_result_code_string(SPI_result)); + + portal = SPI_cursor_open(NULL, plan, NULL, NULL, true); + + for (;;) + { + MemoryContext oldcontext; + uint64 i; + + SPI_cursor_fetch(portal, true, batch_rows); + + if (SPI_tuptable == NULL) + elog(ERROR, "the query did not return a result set"); + + /* + * The descriptor is only available once something has been + * fetched, and the writer needs it before the first row can be + * appended -- so the file is created here rather than before the + * loop. The copy outlives SPI_freetuptable(), which frees the + * descriptor along with the rows it described. + */ + if (tupdesc == NULL) + { + tupdesc = CreateTupleDescCopy(SPI_tuptable->tupdesc); + + rc = routine->open_writer(path, tupdesc, &options, &writer); + if (rc != DL_OK) + dl_error_report(ERROR, rc, "open_writer"); + open_writer = writer; + + rc = dl_arrow_builder_open(tupdesc, &builder); + if (rc != DL_OK) + dl_error_report(ERROR, rc, "open_builder"); + open_builder = builder; + + values = palloc(tupdesc->natts * sizeof(Datum)); + nulls = palloc(tupdesc->natts * sizeof(bool)); + } + + if (SPI_processed == 0) + break; + + oldcontext = MemoryContextSwitchTo(row_context); + + for (i = 0; i < SPI_processed; i++) + { + HeapTuple tuple = SPI_tuptable->vals[i]; + int attno; + + CHECK_FOR_INTERRUPTS(); + + for (attno = 0; attno < tupdesc->natts; attno++) + { + Form_pg_attribute attr = TupleDescAttr(tupdesc, attno); + bool isnull; + Datum value = SPI_getbinval(tuple, SPI_tuptable->tupdesc, + attno + 1, &isnull); + + /* + * The Arrow side runs as C++ and must not allocate, so a + * value that is compressed or stored out of line is + * expanded here, where failing to do so is an ordinary + * error rather than an exception crossing an ABI. + */ + if (!isnull && attr->attlen == -1) + value = PointerGetDatum(PG_DETOAST_DATUM_PACKED(value)); + + values[attno] = value; + nulls[attno] = isnull; + } + + rc = dl_arrow_builder_append(builder, values, nulls, + tupdesc->natts); + if (rc != DL_OK) + dl_error_report(ERROR, rc, "append_row"); + + written++; + } + + MemoryContextSwitchTo(oldcontext); + MemoryContextReset(row_context); + + write_one_batch(writer, builder); + SPI_freetuptable(SPI_tuptable); + } + + SPI_cursor_close(portal); + + /* + * A query that returned nothing still produces a file, with the schema + * and no rows: an empty file is a fact about the query, and a missing + * one would be a fact about this function. + */ + rc = writer->ops->finish(&writer, NULL); + open_writer = NULL; /* consumed, whether or not it succeeded */ + if (rc != DL_OK) + dl_error_report(ERROR, rc, "finish_writer"); + + dl_arrow_builder_close(&builder); + open_builder = NULL; + } + PG_CATCH(); + { + DlArrowBuilder builder = open_builder; + FormatWriter *writer = open_writer; + + if (builder != NULL) + dl_arrow_builder_close(&builder); + if (writer != NULL) + writer->ops->abort(&writer); + + PG_RE_THROW(); + } + PG_END_TRY(); + + SPI_finish(); + + PG_RETURN_INT64(written); +} + +/* + * datalake_parquet_read(path, first_row_group, n_row_groups) -> setof record + * + * The column definition list says what the caller expects the file to hold, and + * is checked against the file's own schema rather than assumed: reading an + * Arrow column as the wrong PostgreSQL type would produce values, just not the + * ones in the file. + * + * The row group arguments are the unit a scan is divided at. Reading 0..0 and + * then 1..1 has to produce exactly what reading the whole file does, which is + * the property a scan spread across segments will depend on. + */ +Datum +datalake_parquet_read(PG_FUNCTION_ARGS) +{ + char *path = text_to_cstring(PG_GETARG_TEXT_PP(0)); + ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo; + const FormatRoutine *routine = parquet_routine(); + FormatReader *volatile open_reader = NULL; + struct ArrowArray *batch = palloc0(sizeof(struct ArrowArray)); + struct ArrowSchema *schema = palloc0(sizeof(struct ArrowSchema)); + Fragment fragment = {0}; + TupleDesc tupdesc; + Tuplestorestate *tupstore; + Datum *values; + bool *nulls; + FormatReader *reader = NULL; + MemoryContext row_context; + DlErrCode rc; + + fragment.path = path; + fragment.first_row_group = PG_GETARG_INT32(1); + fragment.n_row_groups = PG_GETARG_INT32(2); + + InitMaterializedSRF(fcinfo, MAT_SRF_USE_EXPECTED_DESC); + tupdesc = rsinfo->setDesc; + tupstore = rsinfo->setResult; + + values = palloc(tupdesc->natts * sizeof(Datum)); + nulls = palloc(tupdesc->natts * sizeof(bool)); + + /* + * Every text and bytea decoded out of a batch is a copy, and tuplestore + * copies it again. A materialize-mode function is called once, so the + * caller's per-tuple context is not reset until it returns -- without a + * context of its own, reading a large file would hold a second copy of all + * of it until then. + */ + row_context = AllocSetContextCreate(CurrentMemoryContext, + "datalake_parquet_read", + ALLOCSET_DEFAULT_SIZES); + + rc = routine->open_reader(&fragment, NULL, NULL, &reader); + if (rc != DL_OK) + dl_error_report(ERROR, rc, "open_reader"); + open_reader = reader; + + PG_TRY(); + { + for (;;) + { + MemoryContext oldcontext; + bool eof; + int64 row; + int attno; + + CHECK_FOR_INTERRUPTS(); + + rc = reader->ops->next_batch(reader, batch, schema, &eof); + if (rc != DL_OK) + dl_error_report(ERROR, rc, "next_batch"); + if (eof) + break; + + if (schema->n_children != tupdesc->natts) + ereport(ERROR, + (errcode(ERRCODE_DATATYPE_MISMATCH), + errmsg("the file does not have the number of columns the query expects"), + errdetail("The file has %lld, the query expects %d.", + (long long) schema->n_children, + tupdesc->natts))); + + for (attno = 0; attno < tupdesc->natts; attno++) + { + rc = dl_arrow_decode_check(schema->children[attno], + TupleDescAttr(tupdesc, attno)->atttypid); + if (rc != DL_OK) + dl_error_report(ERROR, rc, "check_column"); + } + + oldcontext = MemoryContextSwitchTo(row_context); + + for (row = 0; row < batch->length; row++) + { + CHECK_FOR_INTERRUPTS(); + + for (attno = 0; attno < tupdesc->natts; attno++) + { + rc = dl_arrow_decode_value(batch->children[attno], row, + TupleDescAttr(tupdesc, attno)->atttypid, + TupleDescAttr(tupdesc, attno)->atttypmod, + &values[attno], &nulls[attno]); + if (rc != DL_OK) + dl_error_report(ERROR, rc, "decode_value"); + } + + tuplestore_putvalues(tupstore, tupdesc, values, nulls); + } + + MemoryContextSwitchTo(oldcontext); + MemoryContextReset(row_context); + + /* Releasing the batch releases the columns under it. */ + batch->release(batch); + schema->release(schema); + } + + reader->ops->close(&reader); + open_reader = NULL; + MemoryContextDelete(row_context); + } + PG_CATCH(); + { + FormatReader *failed = open_reader; + + if (batch->release != NULL) + batch->release(batch); + if (schema->release != NULL) + schema->release(schema); + if (failed != NULL) + failed->ops->close(&failed); + + PG_RE_THROW(); + } + PG_END_TRY(); + + return (Datum) 0; +} diff --git a/contrib/datalake_fdw/test/automation/README.md b/contrib/datalake_fdw/test/automation/README.md index f6586ed0e5c..470dba525cb 100644 --- a/contrib/datalake_fdw/test/automation/README.md +++ b/contrib/datalake_fdw/test/automation/README.md @@ -56,13 +56,16 @@ scripts/test/ category runners scripts/utils/ shared shell helpers sqlrepo/smoke/ one directory per category iceberg_am/ DDL, refusals and privileges -- no external service + format_parquet/ a table through a local Parquet file and back ``` -`sqlrepo/smoke/iceberg_am` holds the cases pg_regress runs; the module's -`Makefile` points `--inputdir` here, so `make installcheck` from the module -directory and `make test` from this one run the same cases. They live here -rather than in a `sql/` directory of their own so that there is one place to look -for test material. +Each of these holds cases pg_regress runs, and pg_regress takes one +`--inputdir`, so each category is a run of its own: the module's `Makefile` has +`installcheck` for `iceberg_am` and `installcheck-format-parquet` for the other, +and hangs the second off the first so that asking for `installcheck` gets both. +`make test` from this directory runs the same cases. They live here rather than +in a `sql/` directory of their own so that there is one place to look for test +material. ## What arrives with the metadata engine diff --git a/contrib/datalake_fdw/test/automation/scripts/test/run_smoke_tests.sh b/contrib/datalake_fdw/test/automation/scripts/test/run_smoke_tests.sh index ce707a1fed3..9273c425753 100755 --- a/contrib/datalake_fdw/test/automation/scripts/test/run_smoke_tests.sh +++ b/contrib/datalake_fdw/test/automation/scripts/test/run_smoke_tests.sh @@ -37,9 +37,10 @@ dl_load_config # category:services -- an empty service list means "no external dependency" CATEGORY_SERVICES=" iceberg_am: +format_parquet: " -categories="${CATEGORIES:-iceberg_am}" +categories="${CATEGORIES:-iceberg_am format_parquet}" services_for() { @@ -74,10 +75,18 @@ service_is_available() run_iceberg_am() { # These cases are expected-output cases, so pg_regress runs them; the module - # Makefile already points it at sqlrepo/smoke/iceberg_am. + # Makefile already points it at sqlrepo/smoke/iceberg_am. That target also + # runs format_parquet, so running both categories here runs it twice -- + # which is what "make test CATEGORIES=format_parquet" has to keep working. make -C "$module_dir" USE_PGXS=1 installcheck } +run_format_parquet() +{ + # pg_regress takes one --inputdir, so each category is a run of its own. + make -C "$module_dir" USE_PGXS=1 installcheck-format-parquet +} + failed=0 skipped=0 ran=0 @@ -105,6 +114,7 @@ for category in $categories; do dl_info "RUN $category" case "$category" in iceberg_am) run_iceberg_am ;; + format_parquet) run_format_parquet ;; *) dl_warn "category \"$category\" has no runner"; false ;; esac diff --git a/contrib/datalake_fdw/test/automation/sqlrepo/smoke/format_parquet/expected/parquet_roundtrip.out b/contrib/datalake_fdw/test/automation/sqlrepo/smoke/format_parquet/expected/parquet_roundtrip.out new file mode 100644 index 00000000000..6f649fd1cf4 --- /dev/null +++ b/contrib/datalake_fdw/test/automation/sqlrepo/smoke/format_parquet/expected/parquet_roundtrip.out @@ -0,0 +1,250 @@ +-- Parquet: every type a lake table can store, written to a file and read back, +-- and the row group range that a scan will one day be divided at. +SET client_min_messages = warning; +DROP VIEW IF EXISTS dlparq_roundtrip, dlparq_split; +DROP TABLE IF EXISTS dlparq_src, dlparq_pairs, dlparq_unsupported CASCADE; +CREATE EXTENSION IF NOT EXISTS datalake_fdw; +CREATE EXTENSION IF NOT EXISTS datalake_fdw_test; +RESET client_min_messages; +-- A timestamptz is printed in the session's zone, so without this the output +-- would depend on where the test ran rather than on what the file holds. +SET TimeZone = 'UTC'; +-- Fixed file names rather than a unique one per run: the writer truncates, so a +-- run reuses what the last one left instead of adding to it. Nothing reachable +-- from SQL can remove a file, so unique names would accumulate forever. +\set roundtrip_file '/tmp/datalake_fdw_regress_roundtrip.parquet' +\set split_file '/tmp/datalake_fdw_regress_split.parquet' +\set empty_file '/tmp/datalake_fdw_regress_empty.parquet' +\set batch_file '/tmp/datalake_fdw_regress_batch.parquet' +\set missing_file '/tmp/datalake_fdw_regress_does_not_exist.parquet' +CREATE TABLE dlparq_src ( + c_bool boolean, + c_int2 smallint, + c_int4 integer, + c_int8 bigint, + c_float4 real, + c_float8 double precision, + c_text text, + c_varchar varchar(16), + c_bpchar char(5), + c_bytea bytea, + c_date date, + c_ts timestamp, + c_tstz timestamptz +) DISTRIBUTED RANDOMLY; +-- The dates and timestamps are chosen around both epochs: PostgreSQL counts +-- from 2000-01-01 and Arrow from 1970-01-01, and a value on either side of +-- 1970 is what tells a wrong shift from a right one. A row of nulls is here +-- because a validity bitmap that is never exercised is a bitmap that has not +-- been tested. +INSERT INTO dlparq_src VALUES + (true, 1, 100, 1000, 1.5, 2.5, + 'hello', 'varchar', 'abc', '\x0102'::bytea, + '1970-01-01', '1970-01-01 00:00:00', '1970-01-01 00:00:00+00'), + (false, -2, -200, -2000, -1.5, -2.5, + 'a longer string with 中文', 'x', '', '\x'::bytea, + '2000-01-01', '2000-01-01 12:34:56.789012', '2000-01-01 12:34:56.789012+00'), + (true, 32767, 2147483647, 9223372036854775807, 3.25, 1e300, + '', 'z', 'zzzzz', '\xdeadbeef'::bytea, + '1969-12-31', '1969-12-31 23:59:59.999999', '1969-12-31 23:59:59.999999+00'), + (NULL, NULL, NULL, NULL, NULL, NULL, + NULL, NULL, NULL, NULL, NULL, NULL, NULL); +SELECT datalake_parquet_write(:'roundtrip_file', + 'SELECT * FROM dlparq_src') AS rows_written; + rows_written +-------------- + 4 +(1 row) + +-- The column definition list is what the caller claims the file holds; it is +-- checked against the file's own schema, not assumed. A view so that the list +-- is written once. +CREATE VIEW dlparq_roundtrip AS + SELECT * FROM datalake_parquet_read(:'roundtrip_file') AS t ( + c_bool boolean, + c_int2 smallint, + c_int4 integer, + c_int8 bigint, + c_float4 real, + c_float8 double precision, + c_text text, + c_varchar varchar(16), + c_bpchar char(5), + c_bytea bytea, + c_date date, + c_ts timestamp, + c_tstz timestamptz); +SELECT * FROM dlparq_roundtrip ORDER BY c_int4; + c_bool | c_int2 | c_int4 | c_int8 | c_float4 | c_float8 | c_text | c_varchar | c_bpchar | c_bytea | c_date | c_ts | c_tstz +--------+--------+------------+---------------------+----------+----------+---------------------------+-----------+----------+------------+------------+---------------------------------+------------------------------------- + f | -2 | -200 | -2000 | -1.5 | -2.5 | a longer string with 中文 | x | | \x | 01-01-2000 | Sat Jan 01 12:34:56.789012 2000 | Sat Jan 01 12:34:56.789012 2000 UTC + t | 1 | 100 | 1000 | 1.5 | 2.5 | hello | varchar | abc | \x0102 | 01-01-1970 | Thu Jan 01 00:00:00 1970 | Thu Jan 01 00:00:00 1970 UTC + t | 32767 | 2147483647 | 9223372036854775807 | 3.25 | 1e+300 | | z | zzzzz | \xdeadbeef | 12-31-1969 | Wed Dec 31 23:59:59.999999 1969 | Wed Dec 31 23:59:59.999999 1969 UTC + | | | | | | | | | | | | +(4 rows) + +-- Both directions: one way only says what the file lost, and a file with a row +-- nobody wrote is just as wrong as one missing a row somebody did. +SELECT count(*) AS differences +FROM ((TABLE dlparq_src EXCEPT ALL TABLE dlparq_roundtrip) + UNION ALL + (TABLE dlparq_roundtrip EXCEPT ALL TABLE dlparq_src)) d; + differences +------------- + 0 +(1 row) + +-- The comparison above cannot see the padding of a char(n): bpchar equality +-- ignores trailing spaces, so a value that came back three characters long +-- would still have compared equal to the five it went in as. +SELECT octet_length(c_bpchar) AS bpchar_bytes, + octet_length(c_text) AS text_bytes, + octet_length(c_bytea) AS bytea_bytes +FROM dlparq_roundtrip ORDER BY 1, 2, 3; + bpchar_bytes | text_bytes | bytea_bytes +--------------+------------+------------- + 5 | 0 | 4 + 5 | 5 | 2 + 5 | 27 | 0 + | | +(4 rows) + +-- Row groups are the unit a scan is divided at, so reading the parts has to add +-- up to reading the whole -- no row seen twice, none missed. +CREATE TABLE dlparq_pairs AS + SELECT i AS k, 'v' || i AS v FROM generate_series(1, 6) i + DISTRIBUTED RANDOMLY; +SELECT datalake_parquet_write(:'split_file', + 'SELECT k, v FROM dlparq_pairs ORDER BY k', + 2) AS rows_written; + rows_written +-------------- + 6 +(1 row) + +CREATE VIEW dlparq_split AS + SELECT * FROM datalake_parquet_read(:'split_file') AS t (k int, v text); +SELECT * FROM dlparq_split ORDER BY k; + k | v +---+---- + 1 | v1 + 2 | v2 + 3 | v3 + 4 | v4 + 5 | v5 + 6 | v6 +(6 rows) + +SELECT 0 AS first_row_group, * FROM datalake_parquet_read(:'split_file', 0, 1) + AS t (k int, v text) +UNION ALL +SELECT 1, * FROM datalake_parquet_read(:'split_file', 1, 1) AS t (k int, v text) +UNION ALL +SELECT 2, * FROM datalake_parquet_read(:'split_file', 2, 1) AS t (k int, v text) +ORDER BY 1, 2; + first_row_group | k | v +-----------------+---+---- + 0 | 1 | v1 + 0 | 2 | v2 + 1 | 3 | v3 + 1 | 4 | v4 + 2 | 5 | v5 + 2 | 6 | v6 +(6 rows) + +-- Reading from a row group on is the same as reading each of them. +SELECT count(*) AS rows_from_the_second_on +FROM datalake_parquet_read(:'split_file', 1) AS t (k int, v text); + rows_from_the_second_on +------------------------- + 4 +(1 row) + +-- A query that returns nothing still produces a file: an empty file is a fact +-- about the query, a missing one would be a fact about the writer. +SELECT datalake_parquet_write(:'empty_file', + 'SELECT k, v FROM dlparq_pairs WHERE false') AS rows_written; + rows_written +-------------- + 0 +(1 row) + +SELECT count(*) AS rows_read +FROM datalake_parquet_read(:'empty_file') AS t (k int, v text); + rows_read +----------- + 0 +(1 row) + +-- iceberg.batch_rows is how many rows cross the boundary at a time, and both +-- halves read it. Everything above fits in one batch, which leaves the loops +-- on both sides running exactly once; at two rows a batch the reader iterates +-- and the writer accumulates several batches into the one row group. +SET iceberg.batch_rows = 2; +SELECT count(*) AS rows_in_batches_of_two FROM dlparq_split; + rows_in_batches_of_two +------------------------ + 6 +(1 row) + +SELECT datalake_parquet_write(:'batch_file', + 'SELECT k, v FROM dlparq_pairs ORDER BY k') AS rows_written; + rows_written +-------------- + 6 +(1 row) + +SELECT * FROM datalake_parquet_read(:'batch_file') AS t (k int, v text) ORDER BY k; + k | v +---+---- + 1 | v1 + 2 | v2 + 3 | v3 + 4 | v4 + 5 | v5 + 6 | v6 +(6 rows) + +RESET iceberg.batch_rows; +-- Refusals. Each of these would otherwise be a wrong answer rather than an +-- error: a column silently dropped, a value reinterpreted, a short read. +CREATE TABLE dlparq_unsupported (k int, n numeric) DISTRIBUTED RANDOMLY; +SELECT datalake_parquet_write(:'roundtrip_file', + 'SELECT * FROM dlparq_unsupported'); +ERROR: iceberg: open_writer failed +DETAIL: arrow schema: column "n" has a type that lake tables cannot store yet +SELECT * FROM datalake_parquet_read(:'split_file') AS t (k bigint, v text); +ERROR: iceberg: check_column failed +DETAIL: decode an Arrow column: a column stored as Arrow type "i" cannot be read as bigint +SELECT * FROM datalake_parquet_read(:'split_file') AS t (k int); +ERROR: the file does not have the number of columns the query expects +DETAIL: The file has 2, the query expects 1. +SELECT * FROM datalake_parquet_read(:'split_file', 9, 1) AS t (k int, v text); +ERROR: iceberg: open_reader failed +DETAIL: open a Parquet file: row groups 9..9 were asked for from a file that has 3 +-- A row group range that only overflowed arithmetic would let through: the +-- count is rejected rather than turned into a two-billion-entry list. +SELECT * FROM datalake_parquet_read(:'split_file', 1, 2147483647) AS t (k int, v text); +ERROR: iceberg: open_reader failed +DETAIL: open a Parquet file: row groups 1..2147483647 were asked for from a file that has 3 +SELECT datalake_parquet_write(:'split_file', 'SELECT 1', -1); +ERROR: row group size must be between 0 and 1048576 +SELECT datalake_parquet_write(:'split_file', 'SELECT 1', 2000000000); +ERROR: row group size must be between 0 and 1048576 +-- PostgreSQL's timestamp range runs about 34 years past the last instant Arrow +-- can hold as microseconds from 1970. Writing one of those has to be refused, +-- because the shift would wrap and the value would land 292000 years before the +-- epoch with the write reporting success. +SELECT datalake_parquet_write(:'batch_file', + $$SELECT '294250-01-01 00:00:00'::timestamp$$); +ERROR: iceberg: append_row failed +DETAIL: append a value to an Arrow array: timestamp is too far in the future to be written to a data file (Invalid) +-- Arrow words this one, and its wording is not ours to depend on. +\set VERBOSITY terse +SELECT * FROM datalake_parquet_read(:'missing_file') AS t (k int, v text); +ERROR: iceberg: open_reader failed +\set VERBOSITY default +SET client_min_messages = warning; +DROP VIEW dlparq_roundtrip, dlparq_split; +DROP TABLE dlparq_src, dlparq_pairs, dlparq_unsupported; +RESET client_min_messages; diff --git a/contrib/datalake_fdw/test/automation/sqlrepo/smoke/format_parquet/sql/parquet_roundtrip.sql b/contrib/datalake_fdw/test/automation/sqlrepo/smoke/format_parquet/sql/parquet_roundtrip.sql new file mode 100644 index 00000000000..5b709d20921 --- /dev/null +++ b/contrib/datalake_fdw/test/automation/sqlrepo/smoke/format_parquet/sql/parquet_roundtrip.sql @@ -0,0 +1,174 @@ +-- Parquet: every type a lake table can store, written to a file and read back, +-- and the row group range that a scan will one day be divided at. + +SET client_min_messages = warning; +DROP VIEW IF EXISTS dlparq_roundtrip, dlparq_split; +DROP TABLE IF EXISTS dlparq_src, dlparq_pairs, dlparq_unsupported CASCADE; +CREATE EXTENSION IF NOT EXISTS datalake_fdw; +CREATE EXTENSION IF NOT EXISTS datalake_fdw_test; +RESET client_min_messages; + +-- A timestamptz is printed in the session's zone, so without this the output +-- would depend on where the test ran rather than on what the file holds. +SET TimeZone = 'UTC'; + +-- Fixed file names rather than a unique one per run: the writer truncates, so a +-- run reuses what the last one left instead of adding to it. Nothing reachable +-- from SQL can remove a file, so unique names would accumulate forever. +\set roundtrip_file '/tmp/datalake_fdw_regress_roundtrip.parquet' +\set split_file '/tmp/datalake_fdw_regress_split.parquet' +\set empty_file '/tmp/datalake_fdw_regress_empty.parquet' +\set batch_file '/tmp/datalake_fdw_regress_batch.parquet' +\set missing_file '/tmp/datalake_fdw_regress_does_not_exist.parquet' + +CREATE TABLE dlparq_src ( + c_bool boolean, + c_int2 smallint, + c_int4 integer, + c_int8 bigint, + c_float4 real, + c_float8 double precision, + c_text text, + c_varchar varchar(16), + c_bpchar char(5), + c_bytea bytea, + c_date date, + c_ts timestamp, + c_tstz timestamptz +) DISTRIBUTED RANDOMLY; + +-- The dates and timestamps are chosen around both epochs: PostgreSQL counts +-- from 2000-01-01 and Arrow from 1970-01-01, and a value on either side of +-- 1970 is what tells a wrong shift from a right one. A row of nulls is here +-- because a validity bitmap that is never exercised is a bitmap that has not +-- been tested. +INSERT INTO dlparq_src VALUES + (true, 1, 100, 1000, 1.5, 2.5, + 'hello', 'varchar', 'abc', '\x0102'::bytea, + '1970-01-01', '1970-01-01 00:00:00', '1970-01-01 00:00:00+00'), + (false, -2, -200, -2000, -1.5, -2.5, + 'a longer string with 中文', 'x', '', '\x'::bytea, + '2000-01-01', '2000-01-01 12:34:56.789012', '2000-01-01 12:34:56.789012+00'), + (true, 32767, 2147483647, 9223372036854775807, 3.25, 1e300, + '', 'z', 'zzzzz', '\xdeadbeef'::bytea, + '1969-12-31', '1969-12-31 23:59:59.999999', '1969-12-31 23:59:59.999999+00'), + (NULL, NULL, NULL, NULL, NULL, NULL, + NULL, NULL, NULL, NULL, NULL, NULL, NULL); + +SELECT datalake_parquet_write(:'roundtrip_file', + 'SELECT * FROM dlparq_src') AS rows_written; + +-- The column definition list is what the caller claims the file holds; it is +-- checked against the file's own schema, not assumed. A view so that the list +-- is written once. +CREATE VIEW dlparq_roundtrip AS + SELECT * FROM datalake_parquet_read(:'roundtrip_file') AS t ( + c_bool boolean, + c_int2 smallint, + c_int4 integer, + c_int8 bigint, + c_float4 real, + c_float8 double precision, + c_text text, + c_varchar varchar(16), + c_bpchar char(5), + c_bytea bytea, + c_date date, + c_ts timestamp, + c_tstz timestamptz); + +SELECT * FROM dlparq_roundtrip ORDER BY c_int4; + +-- Both directions: one way only says what the file lost, and a file with a row +-- nobody wrote is just as wrong as one missing a row somebody did. +SELECT count(*) AS differences +FROM ((TABLE dlparq_src EXCEPT ALL TABLE dlparq_roundtrip) + UNION ALL + (TABLE dlparq_roundtrip EXCEPT ALL TABLE dlparq_src)) d; + +-- The comparison above cannot see the padding of a char(n): bpchar equality +-- ignores trailing spaces, so a value that came back three characters long +-- would still have compared equal to the five it went in as. +SELECT octet_length(c_bpchar) AS bpchar_bytes, + octet_length(c_text) AS text_bytes, + octet_length(c_bytea) AS bytea_bytes +FROM dlparq_roundtrip ORDER BY 1, 2, 3; + +-- Row groups are the unit a scan is divided at, so reading the parts has to add +-- up to reading the whole -- no row seen twice, none missed. +CREATE TABLE dlparq_pairs AS + SELECT i AS k, 'v' || i AS v FROM generate_series(1, 6) i + DISTRIBUTED RANDOMLY; + +SELECT datalake_parquet_write(:'split_file', + 'SELECT k, v FROM dlparq_pairs ORDER BY k', + 2) AS rows_written; + +CREATE VIEW dlparq_split AS + SELECT * FROM datalake_parquet_read(:'split_file') AS t (k int, v text); + +SELECT * FROM dlparq_split ORDER BY k; + +SELECT 0 AS first_row_group, * FROM datalake_parquet_read(:'split_file', 0, 1) + AS t (k int, v text) +UNION ALL +SELECT 1, * FROM datalake_parquet_read(:'split_file', 1, 1) AS t (k int, v text) +UNION ALL +SELECT 2, * FROM datalake_parquet_read(:'split_file', 2, 1) AS t (k int, v text) +ORDER BY 1, 2; + +-- Reading from a row group on is the same as reading each of them. +SELECT count(*) AS rows_from_the_second_on +FROM datalake_parquet_read(:'split_file', 1) AS t (k int, v text); + +-- A query that returns nothing still produces a file: an empty file is a fact +-- about the query, a missing one would be a fact about the writer. +SELECT datalake_parquet_write(:'empty_file', + 'SELECT k, v FROM dlparq_pairs WHERE false') AS rows_written; +SELECT count(*) AS rows_read +FROM datalake_parquet_read(:'empty_file') AS t (k int, v text); + +-- iceberg.batch_rows is how many rows cross the boundary at a time, and both +-- halves read it. Everything above fits in one batch, which leaves the loops +-- on both sides running exactly once; at two rows a batch the reader iterates +-- and the writer accumulates several batches into the one row group. +SET iceberg.batch_rows = 2; +SELECT count(*) AS rows_in_batches_of_two FROM dlparq_split; +SELECT datalake_parquet_write(:'batch_file', + 'SELECT k, v FROM dlparq_pairs ORDER BY k') AS rows_written; +SELECT * FROM datalake_parquet_read(:'batch_file') AS t (k int, v text) ORDER BY k; +RESET iceberg.batch_rows; + +-- Refusals. Each of these would otherwise be a wrong answer rather than an +-- error: a column silently dropped, a value reinterpreted, a short read. +CREATE TABLE dlparq_unsupported (k int, n numeric) DISTRIBUTED RANDOMLY; +SELECT datalake_parquet_write(:'roundtrip_file', + 'SELECT * FROM dlparq_unsupported'); + +SELECT * FROM datalake_parquet_read(:'split_file') AS t (k bigint, v text); +SELECT * FROM datalake_parquet_read(:'split_file') AS t (k int); +SELECT * FROM datalake_parquet_read(:'split_file', 9, 1) AS t (k int, v text); + +-- A row group range that only overflowed arithmetic would let through: the +-- count is rejected rather than turned into a two-billion-entry list. +SELECT * FROM datalake_parquet_read(:'split_file', 1, 2147483647) AS t (k int, v text); + +SELECT datalake_parquet_write(:'split_file', 'SELECT 1', -1); +SELECT datalake_parquet_write(:'split_file', 'SELECT 1', 2000000000); + +-- PostgreSQL's timestamp range runs about 34 years past the last instant Arrow +-- can hold as microseconds from 1970. Writing one of those has to be refused, +-- because the shift would wrap and the value would land 292000 years before the +-- epoch with the write reporting success. +SELECT datalake_parquet_write(:'batch_file', + $$SELECT '294250-01-01 00:00:00'::timestamp$$); + +-- Arrow words this one, and its wording is not ours to depend on. +\set VERBOSITY terse +SELECT * FROM datalake_parquet_read(:'missing_file') AS t (k int, v text); +\set VERBOSITY default + +SET client_min_messages = warning; +DROP VIEW dlparq_roundtrip, dlparq_split; +DROP TABLE dlparq_src, dlparq_pairs, dlparq_unsupported; +RESET client_min_messages;