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;